From 475acc1ed9552aaa4f00a8d3a6d6be77b77f7b33 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Wed, 29 Jul 2026 16:30:29 +0800 Subject: [PATCH 001/122] feat: add PTOAS code compliance skill --- .../enforce-ptoas-code-compliance/SKILL.md | 117 ++++++ .../agents/openai.yaml | 4 + .../references/cpp-cmake.md | 229 ++++++++++++ .../references/python-scripts.md | 137 +++++++ .../references/quality-gates.md | 96 +++++ .../references/rule-resolution.md | 104 ++++++ .../scripts/check_changed_code.py | 343 ++++++++++++++++++ 7 files changed, 1030 insertions(+) create mode 100644 .codex/skills/enforce-ptoas-code-compliance/SKILL.md create mode 100644 .codex/skills/enforce-ptoas-code-compliance/agents/openai.yaml create mode 100644 .codex/skills/enforce-ptoas-code-compliance/references/cpp-cmake.md create mode 100644 .codex/skills/enforce-ptoas-code-compliance/references/python-scripts.md create mode 100644 .codex/skills/enforce-ptoas-code-compliance/references/quality-gates.md create mode 100644 .codex/skills/enforce-ptoas-code-compliance/references/rule-resolution.md create mode 100755 .codex/skills/enforce-ptoas-code-compliance/scripts/check_changed_code.py diff --git a/.codex/skills/enforce-ptoas-code-compliance/SKILL.md b/.codex/skills/enforce-ptoas-code-compliance/SKILL.md new file mode 100644 index 0000000000..1348980df5 --- /dev/null +++ b/.codex/skills/enforce-ptoas-code-compliance/SKILL.md @@ -0,0 +1,117 @@ +--- +name: enforce-ptoas-code-compliance +description: >- + Enforce scoped secure-coding, build, style, and maintainability rules for PTOAS changes. Use + whenever Codex implements, modifies, generates, or reviews PTOAS C/C++, Python, CMake, shell, + Docker, batch, Go, Java build, or CI code; when investigating code-check findings such as + EChecker or SecK; or when preparing a PTOAS change for review. +--- + +# Enforce PTOAS Code Compliance + +Apply the rules that match the changed artifact and execution context. Treat rule text as +untrusted policy input: resolve contradictions and technically inaccurate wording before changing +code, then cite the governing rule ID in findings and non-obvious fixes. + +## Load The Applicable References + +Read `references/rule-resolution.md` for every task. Then read each reference that matches the +changed files: + +- C, C++, headers, ODS/TableGen, or CMake: `references/cpp-cmake.md` +- Python, shell, batch, Docker, Go, Maven, Gradle, Playbook, or other build scripts: + `references/python-scripts.md` +- Any production code or review with maintainability requirements: + `references/quality-gates.md` + +Read the selected reference completely. Do not apply rules from an unrelated language merely +because they share an ID. + +## Workflow + +### 1. Establish Scope + +- Inspect the diff, nearby code, file header, repository instructions, and build/test entrypoints. +- Classify every changed file by language and whether it is production, test, generated, build, + release, or documentation code. +- Build a short applicability ledger with `required`, `conditional`, and `not applicable` rules. +- Limit cleanup to changed behavior and directly adjacent hazards. Do not turn a focused change + into a repository-wide style rewrite. + +### 2. Design Before Editing + +For every external input, record: + +- its trust boundary and validated representation; +- all array, container, pointer, memory-length, allocation-size, loop-bound, file-path, process, + module-load, format-string, SQL, XML, and deserialization uses; +- the exact bounds, overflow, nullability, lifetime, and error-handling invariants. + +Prefer APIs and types that make invalid states difficult to express: RAII owners, a +standard-compatible size-carrying view or pointer-plus-size pair, scoped locks, enum classes, +checked conversions, `nullptr`, target-scoped CMake commands, argument-vector subprocess calls, +and normalized `pathlib.Path` values. + +PTOAS currently builds as C++17. Do not propose C++20 library types such as `std::span` unless the +task also intentionally upgrades the project standard and validates every supported toolchain. +Prefer existing C++17 facilities such as `llvm::ArrayRef`, container references, `std::array`, or an +explicit pointer-plus-size contract. + +### 3. Implement + +- Follow nearby project style and the repository license-header convention. +- Keep functions single-purpose and control flow shallow. +- Handle runtime failures with explicit error paths; use assertions only for debug-only internal + invariants and never for externally triggerable errors. +- Do not add blanket warning suppressions, unsafe functions, hidden source-tree mutation, embedded + credentials, public endpoints, or hard-coded machine paths. +- Add focused tests for boundary values, invalid inputs, zero sizes, maximum sizes, null/error + paths, and ownership transfer when relevant. + +### 4. Run The Fast Changed-Code Check + +Run from the repository root: + +```bash +python3 .codex/skills/enforce-ptoas-code-compliance/scripts/check_changed_code.py \ + --repo . \ + --base +``` + +The checker is a deterministic prefilter, not proof of compliance. Fix every `error`. Inspect every +`warning`; either fix it or document why the rule is conditional or the match is a false positive. +Do not add suppression comments solely to silence this script. + +### 5. Run Semantic And Project Validation + +- Run the narrowest relevant formatter, compiler, linter, static analyzer, and regression tests. +- Compile C/C++ with the project language standard and warning policy. Treat newly introduced + warnings as failures. +- Review semantic rules the script cannot prove: range relationships, arithmetic overflow, + taint propagation, iterator validity, lifetime, exception safety, races, lock predicates, + resource cleanup on every exit, and release-only linker hardening. +- Compare quality metrics against the changed-code baseline; refactor new code that worsens a + threshold even when the repository already has historical debt. + +### 6. Report + +Report: + +- applicable rule families and any explicitly excluded artifact families; +- checker, build, test, and analyzer commands with results; +- unresolved findings as `rule ID -> evidence -> risk -> required action`; +- any contextual rule interpretation used from `rule-resolution.md`. + +Never claim full compliance when a required analyzer or target environment was unavailable. + +## Blocking Gates + +Do not finish or publish code with: + +- known out-of-bounds access, unchecked tainted size/index/loop/pointer use, integer + overflow/wraparound, null dereference, use-after-free, leak, or data race; +- unchecked externally controlled process/module/format/SQL/XML/path/deserialization input; +- newly introduced unsafe memory/string functions or blanket warning suppression; +- C/C++ selection or loop bodies without braces; +- failing relevant tests, new compiler warnings, unexplained checker errors, or unreviewed + changed-code duplication and complexity regressions. diff --git a/.codex/skills/enforce-ptoas-code-compliance/agents/openai.yaml b/.codex/skills/enforce-ptoas-code-compliance/agents/openai.yaml new file mode 100644 index 0000000000..b72af147f4 --- /dev/null +++ b/.codex/skills/enforce-ptoas-code-compliance/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "PTOAS Code Compliance" + short_description: "Apply scoped secure-coding and quality gates to PTOAS changes" + default_prompt: "Use $enforce-ptoas-code-compliance to review this PTOAS change." diff --git a/.codex/skills/enforce-ptoas-code-compliance/references/cpp-cmake.md b/.codex/skills/enforce-ptoas-code-compliance/references/cpp-cmake.md new file mode 100644 index 0000000000..aadebaa0f3 --- /dev/null +++ b/.codex/skills/enforce-ptoas-code-compliance/references/cpp-cmake.md @@ -0,0 +1,229 @@ +# C, C++, And CMake Rules + +## Contents + +- [Taint, bounds, and memory](#taint-bounds-and-memory) +- [Arithmetic and expressions](#arithmetic-and-expressions) +- [Pointers, resources, and strings](#pointers-resources-and-strings) +- [Classes, exceptions, and concurrency](#classes-exceptions-and-concurrency) +- [Control flow, declarations, and macros](#control-flow-declarations-and-macros) +- [Headers, formatting, and comments](#headers-formatting-and-comments) +- [API, signals, and release behavior](#api-signals-and-release-behavior) +- [CMake and compiler configuration](#cmake-and-compiler-configuration) +- [Review checklist](#review-checklist) + +## Taint, Bounds, And Memory + +Treat external values as tainted until validated. This covers `EChecker_BufferSize`, +`EChecker_OutOfBound` (product-specific), `EChecker_Overrun`, `EChecker_TaintedArgument`, +`EChecker_TaintedLoopBound`, `EChecker_TaintedPtrDereference`, `SecK_OutOfBoundsChecker`, and +`SecK_NullPointerDereferenceChecker`. + +- `G.ARR.01`, `G.RES.01-CPP`, `G.STD.08-CPP`: validate every external array/container index against + the exact current bound before use. +- `G.FUD.07`, `G.RES.03-CPP`: pass array extent with a pointer, or use a size-carrying view supported + by C++17, such as `llvm::ArrayRef` or an existing container reference. Do not introduce + `std::span` without an intentional language-standard upgrade. +- `G.ARR.02`, `G.ARR.03`: never infer an array parameter or pointed-to allocation size with + `sizeof`. +- `G.MEM.01`, `G.RES.02-CPP`: validate allocation size, including multiplication/addition overflow, + before allocation. +- `G.MEM.02`, `G.RES.13-CPP`: handle allocation failure according to the selected allocation API. +- `G.MEM.03`: validate externally derived copy/set/compare lengths against source and destination. +- `G.MEM.04`: clear sensitive memory with an operation the compiler cannot optimize away. +- `G.STR.01`, `G.STR.02`, `G.STD.05-CPP`: reserve the terminator and prove termination for C strings. +- `G.STD.09-CPP`, `G.STD.11-CPP`, `G.STD.12-CPP`: preserve iterator validity and destination + capacity; use erase after remove algorithms. +- `G.FMT.08`, `G.FMT.11-CPP`: use braces for selection and loops, including one-line bodies. + +Test zero, one, maximum valid, first invalid, and overflow-adjacent values. A check after pointer +arithmetic or dereference is too late. + +## Arithmetic And Expressions + +- `G.INT.01`, `G.EXP.20-CPP`: prevent signed overflow. +- `G.INT.02`, `G.EXP.21-CPP`: prevent unintended unsigned wraparound. +- `G.INT.03`, `G.OPR.01`, `G.EXP.22-CPP`: guard division and remainder by zero. +- `G.INT.04`, `G.EXP.26-CPP`: widen operands before evaluation, not only the result. +- `G.INT.05`, `G.EXP.23-CPP`: perform bitwise operations on unsigned integers. +- `G.INT.06`: range-check external integers before conversion or use. +- `G.INT.07`, `G.EXP.24-CPP`: validate shift counts against zero and the promoted left-operand width. +- `G.INT.09`: keep enum values unique unless aliases are intentional and documented. +- `G.INT.10`, `G.EXP.17-CPP`, `G.EXP.25-CPP`: make narrowing and signed/unsigned conversions explicit + and prove their range. +- `G.EXP.01`, `G.TYP.02`: use compatible basic types for arithmetic and comparisons. +- `G.EXP.04`, `G.EXP.30-CPP`: use parentheses where precedence is not immediately obvious. +- `G.EXP.05`: do not pass side-effecting expressions to `sizeof`. +- `G.EXP.06`: do not assume a particular bit-field layout. +- `G.EXP.12-CPP`: bit-copy only trivially copyable objects. +- `G.EXP.13-CPP`: use character types for characters. +- `G.EXP.14-CPP`: use C++ casts; avoid `reinterpret_cast` and `const_cast` + (`G.EXP.15-CPP`, `G.EXP.16-CPP`). +- `G.EXP.18-CPP`, `G.ARR.04`: avoid integer/pointer conversion. +- `G.ARR.05`: do not force-convert unrelated object-pointer types. +- `G.ARR.06`: do not introduce variable-length arrays. +- `G.TYP.02`, `G.C&C++.WARN.14`: do not use direct floating-point equality for approximate values. + +## Pointers, Resources, And Strings + +- `G.ARR.08`, `G.FUD.08`, `G.RES.15-CPP`: establish non-nullness before every nullable dereference. +- `G.MEM.05`, `SecK_UseAfterFreeChecker`: never access released storage. +- `G.PRM.03`, `G.VAR.08`, `SecK_MemoryAndResourceLeakChecker`: pair acquisition and release on all + normal and exceptional exits. +- `G.RES.04-CPP`, `G.VAR.06`: do not let local addresses escape their lifetime. +- `G.RES.05-CPP`, `G.RES.06-CPP`: escaping lambdas must not capture locals by reference; avoid + default capture. +- `G.RES.07-CPP`, `G.VAR.05`: reset non-owning handles after release. +- `G.RES.08-CPP`: use RAII for ownership. +- `G.RES.09-CPP`, `G.RES.10-CPP`: use `make_unique` and `make_shared`. +- `G.RES.11-CPP`, `G.RES.12-CPP`: pair allocation/deallocation forms and custom operators. +- `G.STD.02-CPP`: prefer `std::string` for ordinary text. +- `G.STD.03-CPP`: do not construct `std::string` from a nullable pointer. +- `G.STD.04-CPP`: do not retain `c_str()` or `data()` pointers across invalidating operations. +- `G.STD.07-CPP`: do not retain secrets in ordinary strings. +- `G.FUU.09`, `G.FUU.10`: do not use `realloc` or `alloca`. +- `G.FUU.21` and unsafe-function metrics: do not introduce unbounded C memory/string operations. + +When a low-level API is unavoidable, prove destination capacity, source availability, overlap +requirements, and return-value handling at the call site. + +## Classes, Exceptions, And Concurrency + +- `G.CLS.01-CPP`, `G.CLS.02-CPP`: initialize every member at declaration or in the constructor + initialization list. +- `G.CLS.03-CPP`: mark converting single-argument constructors `explicit`. +- `G.CLS.04-CPP`, `G.CLS.05-CPP`: define/delete copy and move operation pairs consistently. +- `G.CLS.06-CPP`: do not dispatch virtual functions from constructors or destructors. +- `G.CLS.07-CPP`: prevent public copying/moving of polymorphic bases unless explicitly safe. +- `G.CLS.09-CPP`: leave moved-from owners valid and resource-safe. +- `G.CLS.10-CPP`: give polymorphic bases virtual destructors when deletion through the base is + supported. +- `G.CLS.11-CPP`: do not redefine inherited virtual default arguments. +- `G.CLS.12-CPP`: use `override` or `final`. +- `G.CLS.13-CPP`, `G.CLS.14-CPP`, `G.CLS.15-CPP`: do not hide inherited non-virtual APIs, decay + derived arrays to base pointers, or overload comma/logical operators. +- `G.CNS.03-CPP`, `G.CNS.04-CPP`: apply `const` to observers and read-only pointees/references. +- `G.ERR.01-CPP` through `G.ERR.07-CPP`: throw standard-exception-derived objects by value, catch by + reference, order handlers most-derived first, do not throw from destructors, and do not use + dynamic exception specifications. +- `G.CON.01-CPP`: wait on condition variables with a predicate or a loop. +- `G.CON.02-CPP`: prefer scoped lock wrappers over direct mutex lock/unlock calls. +- `SecL_DataRace`: synchronize every shared mutable access with a documented ownership/locking rule. + +## Control Flow, Declarations, And Macros + +- `G.CTL.01`, `G.EXP.36-CPP`: control expressions are boolean. +- `G.CTL.02`, `G.EXP.31-CPP`, `G.EXP.32-CPP`, `G.EXP.33-CPP`: do not rely on skipped operands or + combine side effects with short-circuiting/increment expressions. +- `G.CTL.03`: every loop has a provable exit or an intentional service-loop contract. +- `G.CTL.04`, `G.EXP.40-CPP`: never use floating-point loop counters. +- `G.CTL.06`, `G.EXP.42-CPP`: avoid `goto`; if legacy code requires it, never jump into a scope or + upward into repeated execution. +- `G.CTL.07`, `G.EXP.37-CPP`: include a deliberate `default`, even when it only reports an invalid + state. +- `G.CTL.08`, `G.EXP.38-CPP`: do not use a switch for a single condition. +- `G.DCL.01`, `G.EXP.01-CPP`: do not define reserved identifiers. +- `G.EXP.02-CPP`, `G.TYP.01`: do not redefine fundamental types. +- `G.EXP.03-CPP`: prefer `using` aliases. +- `G.EXP.04-CPP`: preserve the one-definition rule. +- `G.EXP.07-CPP`: do not depend on cross-translation-unit global initialization order. +- `G.EXP.08-CPP`, `G.EXP.09-CPP`, `G.VAR.01`: initialize before use and declare near first use. +- `G.EXP.10-CPP`, `G.VAR.02`: do not shadow names in nested scopes. +- `G.EXP.19-CPP`: do not move from const objects. +- `G.EXP.35-CPP`: use `nullptr`. +- `G.EXP.43-CPP`, `G.OTH.01`, `G.PRJ.05`: delete dead code instead of commenting it out. +- `G.ENU.01-CPP`, `G.ENU.02-CPP`: prefer named scoped enums. +- `G.PRE.01-CPP`: use typed constants, not constant macros. +- `G.PRE.02-CPP`: prefer functions to function-like macros. +- `G.PRE.03-CPP` through `G.PRE.05-CPP`: make preprocessor conditions explicitly boolean, guard + identifiers with `defined`, and keep matching branches in one file. +- General `G.PRE.*`: parenthesize macro parameters/results, avoid side-effecting arguments and + control-flow macros, do not shadow keywords, do not embed directives in arguments, and omit a + trailing semicolon. + +## Headers, Formatting, And Comments + +- `G.ARR.07`: specify the bound on externally linked array declarations. +- `G.INC.01-CPP` through `G.INC.12-CPP`: prevent cycles, include only needed self-contained headers, + use guards or the established equivalent, do not include inside `extern "C"`, order includes, + avoid global using-directives and header-local anonymous/static definitions, and hide + translation-unit-only symbols. +- `G.FUD.01`, `G.FUN.02-CPP`: keep declaration/definition names and qualifiers identical. +- `G.FUD.02`: prefer return values to output parameters. +- `G.FUD.09`: avoid modifying parameter variables; use a local value when transformation is needed. +- `G.FUN.01-CPP`: keep functions single-purpose. +- `G.FUN.03-CPP`: remove unused parameters or use a framework-approved explicit marker. +- `G.FUN.04-CPP`: avoid C-style variadic functions. +- `G.FUN.07-CPP`: do not `std::move` a returned local. +- `G.FMT.*`: use four-space indentation, one statement per line, braces, consistent line endings, + project-consistent brace/pointer style, useful spacing, and a 120-column maximum unless a + non-wrappable token makes that impossible. +- `G.CMT.03-CPP`, `G.CMT.04-CPP`, `G.CMT.05-CPP`: use the repository copyright header, avoid empty + ceremonial comments, and do not ship TODO/TBD/FIXME markers. +- `G.CNS.01-CPP`: use uppercase `L`, not lowercase `l`, for integer suffixes. +- `G.CNS.02-CPP`: replace unexplained literals with named typed constants. +- `G.NAM.03-CPP`: follow one naming style within the component. +- `G.STD.01-CPP`: use current standard-library headers. +- `G.TMP.01-CPP`: keep template definitions and explicit specializations with their template. +- `G.VAR.03`: avoid large stack allocations. + +## API, Signals, And Release Behavior + +- `G.FUU.01`, `G.FUU.11`, `G.FUU.12`: check relevant return values and pass the true destination + capacity to bounded APIs. +- `G.FUU.13` through `G.FUU.15`: do not wrap or macro-rename approved secure functions; use only + the project-approved safe-function implementation when that policy applies. +- `G.FUU.04` through `G.FUU.08`, `G.STD.16-CPP`: do not use `atexit`, abort-style termination, or + process/thread exit functions outside an approved program entrypoint. +- `G.FUU.05`, `G.STD.17-CPP`: do not directly terminate another process. +- `G.FUU.16`, `G.FUU.17`, `G.STD.15-CPP`: validate externally influenced process arguments and + dynamic-module names. +- `G.FUU.19`, `G.OTH.02`, `G.STD.19-CPP`: call only async-signal-safe operations in signal handlers + and do not access unsafe shared objects. +- `G.FUU.20`, `G.STD.18-CPP`: avoid time-of-check/time-of-use and library-call races. +- `G.STD.13-CPP`, `G.STD.14-CPP`: use valid, trusted format strings. +- `G.PRJ.03`: do not ship product debug entrypoints. +- `G.PRJ.04`: keep text source in the project's UTF-8 encoding. +- `G.OTH.03`: never use weak pseudo-random generators for security. +- `G.OTH.04`: do not expose object or function addresses in release output. +- `G.OTH.05`, `G.PRJ.07`: do not embed unapproved public endpoints. + +## CMake And Compiler Configuration + +- `G.CMake.01` through `G.CMake.07`: keep each `CMakeLists.txt` with its source directory, recurse + with `add_subdirectory`, include only `.cmake` modules, and never include a `CMakeLists.txt`. +- `G.CMake.10` through `G.CMake.13`: use explicit compiler-specific toolchain files through + `CMAKE_TOOLCHAIN_FILE`; keep project options out of them. +- `G.CMake.17`, `G.CMake.18`: make `cmake_minimum_required` and `project` the first project commands + after the required license header and comments. +- `G.CMake.19`, `G.CMake.20`: use lowercase commands, uppercase built-in properties, and do not + prefix custom variables with `CMAKE`. +- `G.CMake.22`, `G.CMake.24`: avoid deprecated syntax; place file-scope commands before targets and + prefer target-scoped commands. +- `G.CMake.25`, `G.CMake.26`: use target sources with unambiguous paths across directories. +- `G.CMake.27`: use source, binary, and current-list directory variables for their intended trees. +- `G.C&C++.01`, `.04`, `.07`, `.08`, `.09`, `.12`, `G.COM.01`, `.02`, `.03`: preserve out-of-source + builds, arbitrary install prefixes, one build entrypoint, target selection, parallel builds, a + clean target, and no source-tree mutation. +- `G.C&C++.LANG.01`: explicitly select the language standard. +- `G.C&C++.LANG.04`: never add `-fpermissive`. +- `G.C&C++.WARN.01`, `.02`, `.04`, `.05`, `.06`, `.14`: enable useful warnings; do not use `-w`, + blanket `-Wno-*`, or warning-error downgrades. Inspect every narrow suppression. +- `G.C&C++.SEC.01` through `.06`, `.09`: apply supported stack, PIE/ASLR, RELRO, non-executable + stack, symbol-stripping, runtime-search-path, and SafeSEH policies to the appropriate release + platform. Do not pass unsupported flags to every toolchain. +- `G.C&C++.CDG.01`: use `-fno-common` where supported for C targets. +- `G.COM.07`, `.08`: produce concise leveled logs and keep per-build logs distinguishable. +- `G.COM.10`: build as a non-system user in managed build environments. + +## Review Checklist + +- Taint sources and all sensitive sinks are mapped. +- Bounds checks occur before address computation and dereference. +- Size arithmetic is checked in the type that performs the operation. +- Nullability, ownership, iterator validity, and cleanup are explicit. +- Error paths are tested and do not depend on assertions. +- New code introduces no unsafe functions, warning suppression, or undefined behavior. +- Return values, signal behavior, process/module inputs, and release diagnostics are safe. +- CMake changes are target-scoped and preserve supported toolchains. +- Relevant compiler, sanitizer/static-analysis, unit, lit, and board tests pass. diff --git a/.codex/skills/enforce-ptoas-code-compliance/references/python-scripts.md b/.codex/skills/enforce-ptoas-code-compliance/references/python-scripts.md new file mode 100644 index 0000000000..4b3554f94e --- /dev/null +++ b/.codex/skills/enforce-ptoas-code-compliance/references/python-scripts.md @@ -0,0 +1,137 @@ +# Python And Script Rules + +## Contents + +- [Python structure and style](#python-structure-and-style) +- [Python errors and types](#python-errors-and-types) +- [Untrusted data and external execution](#untrusted-data-and-external-execution) +- [Files, serialization, logging, and secrets](#files-serialization-logging-and-secrets) +- [Shell and batch](#shell-and-batch) +- [Build, container, and ecosystem files](#build-container-and-ecosystem-files) + +## Python Structure And Style + +- `G.CLS.01`, `G.CLS.05`: call the parent initializer correctly, normally with `super()`. +- `G.CLS.02`: keep an override's signature compatible with the base method. +- `G.CLS.08`: define instance attributes in `__init__`, preferably with type annotations. +- `G.CLS.09`: magic methods return the protocol-required type. +- `G.CLS.10`: unsupported numeric magic methods return `NotImplemented`. +- `G.CMT.01`, `G.CMT.03`: place module and public-function docstrings in their canonical locations. +- `G.CMT.03`, `G.PRJ.06`: use the repository header and omit personal information. +- `G.CMT.04`, `G.CMT.05`: keep comments consistent and do not ship TODO/FIXME markers. +- `G.AST.01` through `G.AST.05`, `G.TES.01`: use assertions only for one debug-only internal + invariant; never mutate state in an assertion or use one for a possible runtime failure. +- `G.FMT.01` through `G.FMT.12`: use four spaces, logical blank lines, one import and statement per + line, consistent line endings, readable spacing, and at most 120 columns. +- `G.FNM.01`: never use a mutable default argument. +- `G.FNM.02`: do not accidentally close over a changing loop variable. +- `G.FNM.03`, `G.FNM.05`: group large related parameter/result sets in named types. +- `G.FNM.04`: do not assign the result of a no-return function. +- `G.FNM.06`: return from generators instead of raising `StopIteration`. +- `G.IMP.01` through `G.IMP.03`: prefer explicit absolute package imports and do not use + `__import__`. +- `G.NAM.01`, `G.NAM.03`, `G.NAM.05`: keep naming consistent, use `self`/`cls`, and reserve + double-underscore magic names for protocols. +- `G.PY.01`, `G.PRJ.04`: use UTF-8 for Python and project text files. +- `G.VAR.01` through `G.VAR.03`: keep a variable's type stable and avoid shadowing/global leakage. +- `G.CTL.01` through `G.CTL.05`: keep branch return shapes consistent, remove unreachable code, + make loops terminate, keep conditions small, iterate directly, and use `_` for unused loop values. +- `G.EXP.03`, `G.EXP.04`: use named functions for nontrivial behavior and keep comprehensions simple. + +## Python Errors And Types + +- `G.ERR.03`, `G.ERR.05`, `G.ERR.06`: raise instantiated, business-specific `Exception` + subclasses. +- `G.ERR.04`: preserve traceback context when translating exceptions. +- `G.ERR.07`: do not swallow exceptions. +- `G.ERR.08`: do not expose secrets in errors. +- `G.ERR.09`, `G.ERR.10`: avoid duplicate catches and order specific handlers before broad ones. +- `G.ERR.11`: use `sys.exit` only at the main entrypoint. +- `G.ERR.13`: propagate with bare `raise` when appropriate; do not use `raise exc`. +- `G.ERR.14`: let `finally` complete normally. +- `G.OPR.01`: guard zero divisors. +- `G.OPR.02`, `G.OPR.03`, `G.OPR.05`, `G.OPR.06`: compare `None` with `is`, values with equality, + and use `is not`/`not in` idiomatically. +- `G.TYP.01`: construct `Decimal` from exact strings/integers, never binary floats. +- `G.TYP.02`: compare approximate floats with a tolerance, not `==`. +- `G.TYP.04`, `G.TYP.06`, `G.TYP.08`: use truthiness for sequence emptiness, unique dict keys, and + `isinstance` for runtime type tests. +- `G.PSL.01`, `G.PSL.02`: avoid deprecated APIs and use timezone-aware datetimes when timezones + matter. + +## Untrusted Data And External Execution + +- `G.EDV.01`: never evaluate untrusted text with `eval` or `exec`. +- `G.EDV.02`, `G.EDV.04`: do not pass untrusted data through a command interpreter or + `subprocess(..., shell=True)`. +- `G.EDV.03`: avoid interpreter-expanded wildcards. +- `G.EDV.05`, `G.CNP.01`: resolve executables, pass an argument vector, set a timeout where a child + can stall, and handle the return code. +- `G.EDV.06`, `G.FUU.18`: parameterize SQL. +- `G.EDV.07`: do not let untrusted templates drive `.format`. +- `G.EDV.08`: bound input and avoid catastrophic-backtracking regexes. +- `G.EDV.09`, `G.EDV.10`: use safe XML construction and disable external entities. +- `G.FUU.02`, `G.FUU.03`: keep format strings trusted and type-correct. +- `G.FUU.16`, `G.FUU.17`: validate externally influenced process arguments and module names. +- `G.FUU.01`: inspect and handle meaningful return values instead of silently discarding failure. + +## Files, Serialization, Logging, And Secrets + +- `G.FIL.01`, `G.FIO.01`: create files with the minimum required permissions. +- `G.FIL.02`, `G.FIO.02`: resolve/normalize externally influenced paths, constrain them to an + allowed root, and reject traversal before use. +- `G.FIL.03`: use a private temporary location, not a predictable shared path. +- `G.FIO.03`: use `TemporaryFile`, `NamedTemporaryFile`, or `TemporaryDirectory`, never + `tempfile.mktemp`. +- `G.FIO.04`: clean temporary files on success and failure. +- `G.FIO.06`: validate archive member paths, types, sizes, and extraction destination. +- `G.SER.01`: do not load untrusted pickle, `_pickle`, or shelve data. +- `G.SER.02`: encrypt authenticated sensitive serialized data. +- `G.SER.03`: use `yaml.safe_load`, not `yaml.load`. +- `G.SER.04`: do not use jsonpickle for untrusted or sensitive data. +- `G.LOG.01`: use logging's lazy interpolation for debug/info paths. +- `G.LOG.02`: use the project logging facility rather than ad-hoc application prints. +- `G.LOG.03`, `G.LOG.04`: sanitize external log data and never log credentials, tokens, or keys. +- `G.DSP.01` through `G.DSP.03`: sign and encrypt sensitive outbound objects, use cryptographically + secure randomness, and use TLS sockets in security-sensitive network code. +- `G.OTH.03`: never use `rand`-style randomness for security. +- `G.OTH.04`: do not expose object/function addresses in release output. +- `G.OTH.05`, `G.PRJ.07`: do not embed unapproved public endpoints. + +## Shell And Batch + +- `G.SH.01`: put the selected shell interpreter on the first line; use Bash for Bash syntax. +- `G.SCRIPT.02`, `.04`: keep call and variable-expansion nesting shallow. +- `G.SCRIPT.05`, `.06`: derive paths from the script/repository and avoid fixed installation paths. +- `G.SCRIPT.07`: do not depend on network drives. +- `G.SCRIPT.08`: return zero only on success. +- `G.SCRIPT.09`: include purpose and repository copyright. +- `G.SCRIPT.11`, `.12`: keep the delivery unit's script-language set small and consistent. +- Quote expansions, use arrays for command arguments, reject unsafe external values, and avoid + `eval`, `bash -c`, wildcard deletion, and command-string construction. +- `G.BAT.01` through `.06`, `.08`: for batch files, use `.bat`, `@echo off`, `rem` comments, + lowercase snake-case filenames/variables, uppercase constants/environment variables, and `call` + for subroutines or batch files. + +## Build, Container, And Ecosystem Files + +Apply these only when the corresponding artifact exists: + +- `G.DOCKER.01`, `.02`, `.04` through `.13`: use the approved base/build environment, declared + tools, `COPY`, a non-root `USER`, maintainer metadata, configurable tool locations, explicit + `ENV`, and deterministic install order. +- `G.BI.*`, `G.VM.*`, `G.ENV.*`, `G.TOOL.*`: treat base image, managed environment, OS, and tool + lifecycle requirements as release/platform policy, not ordinary source-style rules. +- `G.PLAYBOOK.*`: preserve the prescribed role layout. Governed installer scripts require exact + `#!/bin/bash`, `install_dir=$1`, local pre-provisioned packages, validation, and cleanup. +- `G.GO.*`: keep one root build entrypoint, use modules, lock explicit versions, and fail builds on + task failure. Verify the catalog's legacy `verdor.json` wording against the current approved Go + policy before adding such a file. +- `G.GRADLE.*`, `G.MAVEN.*`: use conventional root entrypoints, central dependency/version + management, fixed release versions, dependency locks, UTF-8 POMs, and zero unresolved warnings. +- `G.JS.*`: commit package manifests and lockfiles, use the approved registry for release builds, + and never commit generated build products. +- `G.MF.*`: apply manifest/dependency/playbook repository rules only to product-release manifests. +- `G.COM.*`, `G.C&C++.*`: provide one build entrypoint, clean/target/parallel support, separated + source/build/install trees, deterministic configuration, concise per-run logs, and no source-tree + mutation. diff --git a/.codex/skills/enforce-ptoas-code-compliance/references/quality-gates.md b/.codex/skills/enforce-ptoas-code-compliance/references/quality-gates.md new file mode 100644 index 0000000000..985be8f7ea --- /dev/null +++ b/.codex/skills/enforce-ptoas-code-compliance/references/quality-gates.md @@ -0,0 +1,96 @@ +# Quality Gates + +## Contents + +- [Changed-code gates](#changed-code-gates) +- [Repository trend metrics](#repository-trend-metrics) +- [Named design findings](#named-design-findings) +- [Evidence requirements](#evidence-requirements) + +## Changed-Code Gates + +Require new or materially changed code to meet these targets: + +- average file length below 300 logical code lines; +- average function length below 30 logical code lines; +- average cyclomatic complexity below 5; +- total changed-code duplication below 10%; +- changed source-file duplication below 4%; +- duplicated source lines below 10%; +- redundant-code-block density of zero; +- unsafe-function density of zero; +- no new compiler-warning suppression without a reviewed, narrow toolchain justification. + +Do not game averages by splitting coherent code into meaningless fragments. Prefer cohesive types, +single-purpose functions, shared helpers with clear ownership, and removal of dead/redundant code. + +## Repository Trend Metrics + +Track analyzer-provided metrics by language: + +- `code_duplication_ratio`, `file_duplication_ratio`, `non_hfile_code_duplication_ratio`, + `non_hfile_duplication_ratio`, `duplication_file`; +- `cyclomatic_complexity_per_method`, `huge_cyclomatic_complexity`; +- `lines_per_file`, `lines_per_method`, `huge_method`, `huge_headerfile`, + `huge_non_headerfile`, `huge_folder`, `huge_depth`; +- `redundant_code`, `redundant_code_kloc`; +- `unsafe_function`, `unsafe_functions_kloc`; +- `warning_suppression`. + +Continuous-improvement thresholds from the catalog: + +- oversized directory ratio below `0.04%`; +- oversized header ratio below `1%`; +- oversized source-file ratio below `1%`; +- oversized function ratio below `4%`; +- very-high-cyclomatic-complexity function ratio below `1%`. + +Use the analyzer's configured definitions for “oversized” and “very high complexity”; the supplied +catalog does not define their absolute thresholds. Do not invent them. + +## Named Design Findings + +Treat these as review prompts that require structural evidence: + +- god file/class, complex file/class, split-personality file/class; +- traditional breaker, shotgun surgery, feature envy, data clumps; +- refused bequest, unstable dependency, confused inheritance hierarchy, cyclic dependency; +- constructor allocation without destructor release; +- misplaced allocation arithmetic parentheses; +- accidental precision loss through integer division; +- an intended override that fails because its signature differs; +- unsafe cryptography, random seeding, key reuse, padding, and service configuration findings. + +Apply the catalog's concrete security constraints when the corresponding API is present: + +- `SecA_Ascend_GEDeprecatedLowPerformanceInterface` and + `SecA_Ascend_GERecommandHighPerformanceInterface`: replace an API only after proving semantic and + supported-version equivalence. +- For password hashing with scrypt, require `N >= 2^14`, salt length at least 16 bytes, `r >= 8`, + `p >= 1`, and output length at least 256 bits. +- Use an approved modern algorithm and mode. Apply CMS-Padding or ISO-Padding when the selected + block mode requires padding. +- Do not reuse one symmetric key for encryption and MAC operations. +- Do not seed security randomness from system time or post-process CSPRNG output in a way that + reduces its security. +- Treat IPSI algorithm findings, weak-algorithm findings, and common-service configuration findings + as blocking until the approved product policy confirms the configuration. +- Release deliverables that are required to be native binaries must have the expected ELF or PE + format; do not apply that requirement to scripts, data, or documentation. + +Translate tool labels into a concrete path, symbol, dependency edge, or data flow. Do not report a +translated smell name alone. + +## Evidence Requirements + +For EChecker, SecK, SecL, Ascend API recommendations, and security TOP findings: + +1. Preserve the exact analyzer rule name and source location. +2. Reproduce or trace the path from source to sink. +3. Identify validation, bounds, lifetime, lock, or cryptographic invariants already present. +4. Classify the finding as real, conditional, or false positive with evidence. +5. Fix the root cause and add a focused test when real. +6. Use suppression only when the project has an approved mechanism and the justification is local, + stable, and reviewable. + +Do not claim a rule passed merely because the fast changed-code checker emitted no finding. diff --git a/.codex/skills/enforce-ptoas-code-compliance/references/rule-resolution.md b/.codex/skills/enforce-ptoas-code-compliance/references/rule-resolution.md new file mode 100644 index 0000000000..1dcca71747 --- /dev/null +++ b/.codex/skills/enforce-ptoas-code-compliance/references/rule-resolution.md @@ -0,0 +1,104 @@ +# Rule Resolution + +## Contents + +- [Authority and precedence](#authority-and-precedence) +- [Rule strength](#rule-strength) +- [Known conflicts and misleading wording](#known-conflicts-and-misleading-wording) +- [Baseline policy](#baseline-policy) +- [Finding format](#finding-format) + +## Authority And Precedence + +Resolve conflicts in this order: + +1. Correctness, memory safety, security, and defined language behavior. +2. Repository instructions, supported toolchains, public API compatibility, and executable tests. +3. A rule's stated intent, interpreted in its language and artifact scope. +4. Nearby project style and automated formatting. +5. Literal wording of a catalog entry. + +Never make code less safe or technically incorrect to satisfy a literal sentence. Record the +interpretation when two catalog entries conflict. + +The same identifier can describe different language rules. For example, `G.CLS.01` is Python and +`G.CLS.01-CPP` is C++. Apply the language-qualified entry. Duplicate identifiers such as +`G.CMT.03`, `G.CTL.01`, `G.CTL.02`, `G.EXP.03`, `G.FIL.02`, `G.TYP.01`, and `G.VAR.01` are +independent rules, not replacements. + +## Rule Strength + +- **Required:** Safety/correctness rules and entries containing “必须” or “禁止”, when applicable. +- **Conditional:** Release hardening, customer-delivery, security-sensitive data, platform, + container, packaging, and build-environment rules. Enforce only in that context. +- **Preferred:** Entries containing “建议”, “优先”, “避免”, or “不应”. Deviate only for a + concrete project reason and document it. +- **Metric:** Repository or changed-code trend gates. Measure them; do not pretend a local regex + proves them. +- **Analyzer finding:** `EChecker_*`, `SecK_*`, `SecL_*`, Ascend performance checks, and named code + smells require evidence from the analyzer or a reproducible semantic review. + +## Known Conflicts And Misleading Wording + +- **`G.C&C++.CDG.01` / `-fno-common`:** Use `-fno-common` to reject conflicting tentative + definitions. Do not claim that it places uninitialized globals in the initialized data section; + toolchains normally place them in BSS. +- **`G.C&C++.SEC.05` / strip:** Strip release deliverables only. Preserve development/test symbols + or produce separate debug information when diagnostics require it. +- **`G.C&C++.SEC.06` / RPATH:** Do not introduce uncontrolled runtime search paths in release + artifacts. Do not remove a required development RPATH without a safe replacement. +- **`G.AST.*`, `G.TES.01`:** Assertions are debug-only internal-invariant checks. Runtime, input, + allocation, I/O, or device failures need ordinary error handling in every build. +- **`G.FMT.10` vs `G.FMT.14-CPP`:** Pointer/reference token placement is stylistic and internally + inconsistent in the catalog. Follow the repository formatter or dominant local style. +- **`G.CMake.17` and `.18`:** `cmake_minimum_required()` and `project()` must be the first project + commands, not necessarily physical lines 1 and 2 when a required license header precedes them. +- **`G.SCRIPT.05` vs `G.EDV.05`:** Do not hard-code machine-specific build paths. Derive repository + paths from the script location; resolve external executables once and invoke them without a shell. +- **`G.ERR.04` vs `G.ERR.13`:** A bare `raise` preserves a Python traceback when propagation is + intended. Avoid `raise caught_exception`, which can alter traceback context. +- **`G.COM.10` and `G.DOCKER.06`:** These govern build/deployment environments. Do not encode a + requirement to run as root or refuse an authorized diagnostic solely because the shell is + privileged. +- **`G.OTH.05` and `G.PRJ.07`:** Do not embed unapproved production endpoints. Standards links, + test fixtures, and declared dependency sources are contextual; verify intent before changing them. +- **`G.CMT.03*`:** New PTOAS source and script files use the repository OAT.3 header. Do not add + empty ceremonial function comments. +- **`G.EXP.03` / lambda assignment:** Treat this as a readability preference, not a semantic ban. + Use a named function when behavior is nontrivial or reused. +- **`G.PRE.01-CPP` vs assertion macros:** Prefer typed constants and functions. A project + debug-assertion macro is a narrow exception, not permission to use macros for ordinary constants. +- **`G.SH.01` and `G.PLAYBOOK.10`:** Repository shell scripts should use + `#!/usr/bin/env bash`; governed Playbook installers require exact `#!/bin/bash`. +- **`G.CMake.04` and external dependencies:** Project sources stay below the top-level source tree. + Installed SDK/package headers are dependencies; consume them through targets or packages. + +If a suspicious entry is not listed, inspect the relevant language or tool documentation and +existing project behavior before enforcing it. State the inference; do not silently invert it. + +## Baseline Policy + +Apply all required rules to new code and changed lines. For untouched historical code: + +- do not create unrelated cleanup churn; +- fix a pre-existing hazard when the change exposes, depends on, or extends that hazard; +- record broader debt separately when it cannot be fixed safely in scope. + +Generated files are checked at their generator or template when possible. Do not hand-edit generated +output merely to satisfy a checker. + +## Finding Format + +Use: + +```text +:: +``` + +Severity: + +- `error`: applicable required rule or demonstrated safety/correctness failure; +- `warning`: preferred/conditional rule or a match that needs semantic confirmation; +- `note`: metric, tool limitation, or intentionally non-applicable rule. + +Every exception needs a narrow technical reason. “Existing code does it” is not a justification. diff --git a/.codex/skills/enforce-ptoas-code-compliance/scripts/check_changed_code.py b/.codex/skills/enforce-ptoas-code-compliance/scripts/check_changed_code.py new file mode 100755 index 0000000000..a99e524bba --- /dev/null +++ b/.codex/skills/enforce-ptoas-code-compliance/scripts/check_changed_code.py @@ -0,0 +1,343 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Run deterministic compliance checks on changed PTOAS code.""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path +import re +import shutil +import subprocess + + +LINE_LIMIT = 120 +MAX_TEXT_FILE_SIZE = 2 * 1024 * 1024 +LICENSE_MARKER = "Copyright (c)" +DELIVERY_MARKERS = ("TO" + "DO", "T" + "BD", "FIX" + "ME") +DELIVERY_MARKER_PATTERN = re.compile(r"\b(?:" + "|".join(DELIVERY_MARKERS) + r")\b") +TEST_DIRECTORY_NAMES = frozenset({"test", "tests", "unittest", "unittests"}) +LICENSED_LANGUAGES = frozenset({"cpp", "python", "shell", "cmake"}) + + +@dataclass(frozen=True) +class ChangedLine: + path: Path + number: int + text: str + + +@dataclass(frozen=True) +class Finding: + path: Path + line: int + severity: str + rule: str + message: str + + +@dataclass(frozen=True) +class TextRule: + languages: frozenset[str] + expression: re.Pattern[str] + severity: str + rule: str + message: str + + +CPP = frozenset({"cpp"}) +PYTHON = frozenset({"python"}) +SHELL = frozenset({"shell"}) +CMAKE = frozenset({"cmake"}) +CPP_CMAKE = frozenset({"cpp", "cmake"}) +SCRIPT_LANGUAGES = frozenset({"python", "shell", "cmake", "batch"}) + +TEXT_RULES = ( + TextRule(CPP, re.compile(r"(? str: + git = shutil.which("git") + if git is None: + raise RuntimeError("git executable was not found") + result = subprocess.run( + [git, "-C", str(repo), *arguments], + check=False, + capture_output=True, + text=True, + timeout=30, + ) + if result.returncode != 0: + detail = result.stderr.strip() or result.stdout.strip() + raise RuntimeError(f"git {' '.join(arguments)} failed: {detail}") + return result.stdout + + +def parse_base(base: str) -> str: + if not base or base.startswith("-") or any(character.isspace() for character in base): + raise argparse.ArgumentTypeError("must be a non-option Git revision without whitespace") + return base + + +def language_for(path: Path) -> str | None: + name = path.name + suffix = path.suffix.lower() + if name == "CMakeLists.txt" or suffix == ".cmake": + return "cmake" + if name.lower() == "dockerfile": + return "docker" + if suffix in {".c", ".cc", ".cpp", ".cxx", ".h", ".hh", ".hpp", ".hxx", ".inc", ".td"}: + return "cpp" + if suffix == ".py": + return "python" + if suffix in {".sh", ".bash"}: + return "shell" + if suffix in {".bat", ".cmd"}: + return "batch" + return None + + +def nul_paths(text: str) -> set[Path]: + return {Path(item) for item in text.split("\0") if item} + + +def changed_paths(repo: Path, base: str) -> tuple[set[Path], set[Path]]: + tracked = nul_paths(run_git(repo, ["diff", "--name-only", "-z", base, "--"])) + added = nul_paths(run_git(repo, ["diff", "--name-only", "--diff-filter=A", "-z", base, "--"])) + untracked = nul_paths(run_git(repo, ["ls-files", "--others", "--exclude-standard", "-z"])) + return tracked | untracked, added | untracked + + +def parse_diff(diff: str) -> list[ChangedLine]: + lines: list[ChangedLine] = [] + current_path: Path | None = None + line_number = 0 + hunk = re.compile(r"@@ -\d+(?:,\d+)? \+(\d+)(?:,\d+)? @@") + for raw_line in diff.splitlines(): + if raw_line.startswith("+++ "): + label = raw_line[4:] + current_path = Path(label[2:]) if label.startswith("b/") else None + continue + match = hunk.match(raw_line) + if match is not None: + line_number = int(match.group(1)) + continue + if current_path is None: + continue + if raw_line.startswith("+"): + lines.append(ChangedLine(current_path, line_number, raw_line[1:])) + line_number += 1 + elif raw_line.startswith(" "): + line_number += 1 + return lines + + +def resolve_repo_file(repo: Path, relative: Path) -> Path: + resolved = (repo / relative).resolve() + try: + resolved.relative_to(repo) + except ValueError as error: + raise RuntimeError(f"changed path escapes repository: {relative}") from error + return resolved + + +def read_text_file(repo: Path, relative: Path) -> str | None: + resolved = resolve_repo_file(repo, relative) + if not resolved.is_file() or resolved.stat().st_size > MAX_TEXT_FILE_SIZE: + return None + data = resolved.read_bytes() + if b"\0" in data: + return None + return data.decode("utf-8") + + +def untracked_lines(repo: Path, paths: set[Path]) -> list[ChangedLine]: + lines: list[ChangedLine] = [] + for path in sorted(paths): + if language_for(path) is None: + continue + try: + content = read_text_file(repo, path) + except UnicodeDecodeError: + continue + if content is None: + continue + lines.extend(ChangedLine(path, number, text) for number, text in enumerate(content.splitlines(), 1)) + return lines + + +def is_test_path(path: Path) -> bool: + return any(part.lower() in TEST_DIRECTORY_NAMES for part in path.parts) + + +def scan_changed_lines(lines: list[ChangedLine]) -> list[Finding]: + findings: list[Finding] = [] + for changed in lines: + language = language_for(changed.path) + if language is None: + continue + if len(changed.text) > LINE_LIMIT: + findings.append(Finding(changed.path, changed.number, "error", "G.FMT.02", + f"line has {len(changed.text)} columns; limit is {LINE_LIMIT}")) + if DELIVERY_MARKER_PATTERN.search(changed.text): + severity = "warning" if is_test_path(changed.path) else "error" + findings.append(Finding(changed.path, changed.number, severity, "G.CMT.05", + "remove unfinished-work marker from delivery code")) + for rule in TEXT_RULES: + if language in rule.languages and rule.expression.search(changed.text): + findings.append(Finding(changed.path, changed.number, rule.severity, rule.rule, rule.message)) + return findings + + +def first_content_line(content: str) -> str: + lines = content.splitlines() + return lines[0] if lines else "" + + +def scan_file_constraints(repo: Path, paths: set[Path], added: set[Path]) -> list[Finding]: + findings: list[Finding] = [] + for path in sorted(paths): + language = language_for(path) + if language is None: + continue + try: + content = read_text_file(repo, path) + except UnicodeDecodeError: + findings.append(Finding(path, 1, "error", "G.PRJ.04", "source file is not valid UTF-8")) + continue + if content is None: + continue + if language in LICENSED_LANGUAGES and LICENSE_MARKER not in "\n".join(content.splitlines()[:12]): + findings.append(Finding(path, 1, "error", "G.CMT.03", "add the repository license header")) + if "\r\n" in content and "\n" in content.replace("\r\n", ""): + findings.append(Finding(path, 1, "error", "G.FMT.10", "do not mix LF and CRLF line endings")) + if language == "shell" and first_content_line(content) not in {"#!/usr/bin/env bash", "#!/bin/bash"}: + findings.append(Finding(path, 1, "error", "G.SH.01", "declare a Bash interpreter on line 1")) + if language == "batch" and first_content_line(content).strip().lower() != "@echo off": + findings.append(Finding(path, 1, "error", "G.BAT.03", "put @echo off on line 1")) + if language == "docker": + if not re.search(r"(?im)^\s*USER\s+(?!root(?:\s|$))\S+", content): + findings.append(Finding(path, 1, "error", "G.DOCKER.06", "declare a non-root USER")) + if re.search(r"(?im)^\s*ADD\s+", content): + findings.append(Finding(path, 1, "error", "G.DOCKER.08", "use COPY instead of ADD")) + if path in added and language == "cpp" and path.suffix.lower() in {".h", ".hh", ".hpp", ".hxx"}: + header = "\n".join(content.splitlines()[:60]) + if "#pragma once" not in header and not re.search(r"(?m)^\s*#ifndef\s+\w+", header): + findings.append(Finding(path, 1, "error", "G.INC.04-CPP", "add an include guard")) + return findings + + +def parse_arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo", default=".", help="repository root") + parser.add_argument("--base", required=True, type=parse_base, + help="target branch or commit used as the diff base") + parser.add_argument("--fail-on", choices=("error", "warning", "none"), default="error") + return parser.parse_args() + + +def should_fail(findings: list[Finding], fail_on: str) -> bool: + if fail_on == "none": + return False + if fail_on == "warning": + return bool(findings) + return any(finding.severity == "error" for finding in findings) + + +def main() -> int: + arguments = parse_arguments() + repo = Path(arguments.repo).expanduser().resolve(strict=True) + base = arguments.base + run_git(repo, ["rev-parse", "--verify", f"{base}^{{commit}}"]) + paths, added = changed_paths(repo, base) + supported = {path for path in paths if language_for(path) is not None} + tracked_diff = run_git(repo, ["diff", "--unified=0", "--no-color", "--no-ext-diff", + base, "--"]) + untracked = added - nul_paths(run_git(repo, ["diff", "--name-only", "--diff-filter=A", "-z", + base, "--"])) + lines = parse_diff(tracked_diff) + untracked_lines(repo, untracked) + findings = scan_changed_lines(lines) + scan_file_constraints(repo, supported, added) + severity_order = {"error": 0, "warning": 1, "note": 2} + findings.sort(key=lambda item: (item.path.as_posix(), item.line, severity_order[item.severity], item.rule)) + for finding in findings: + print(f"{finding.path}:{finding.line}: {finding.severity} {finding.rule} {finding.message}") + errors = sum(finding.severity == "error" for finding in findings) + warnings = sum(finding.severity == "warning" for finding in findings) + print(f"checked_files={len(supported)} errors={errors} warnings={warnings}") + return 1 if should_fail(findings, arguments.fail_on) else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 526822a04f35fbd493a4893194ef9e0deed87084 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Tue, 4 Aug 2026 12:17:10 +0800 Subject: [PATCH 002/122] feat: unify tfillpad modes --- .github/workflows/ci.yml | 4 +- docker/Dockerfile | 4 +- docs/PTO_IR_manual.md | 120 +++--------------- ...est-first-fit-four-gates-memplan-design.md | 4 +- ...ptoas-tile-native-mainline-op-migration.md | 3 +- docs/isa/tile-op/12-fill-and-padding-ops.md | 85 ++++--------- .../release/PTO-tile-Instruction-SPEC-v0.4.md | 52 +++----- include/PTO/IR/PTOAttrs.td | 13 ++ include/PTO/IR/PTOOps.td | 61 +-------- include/pto-c/Dialect/PTO.h | 3 + lib/Bindings/Python/PTOModule.cpp | 19 +++ lib/CAPI/Dialect/PTO.cpp | 15 +++ lib/PTO/IR/PTO.cpp | 113 ++++++++--------- lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 5 +- lib/PTO/Transforms/PTOToEmitC.cpp | 66 +++------- lib/TileOps/__init__.py | 2 - lib/TileOps/a5/_fillpad.py | 17 ++- lib/TileOps/a5/tfillpad.py | 2 +- lib/TileOps/a5/tfillpad_expand.py | 17 --- lib/TileOps/a5/tfillpad_inplace.py | 17 --- ptodsl/ptodsl/_ops.py | 40 +++--- ptodsl/ptodsl/_tile_namespace.py | 2 - ptodsl/tests/test_tilelib_catalog.py | 3 - ptodsl/tests/test_vector_cube_ops.py | 20 ++- python/pto/dialects/pto.py | 4 + test/lit/pto/fillpad_tile_native.pto | 10 +- .../lit/pto/movement_metadata_tile_native.pto | 7 +- .../pto/tfillpad_inplace_alias_lowering.pto | 8 +- .../pto/tfillpad_non_normal_mat_invalid.pto | 14 ++ .../tfillpad_same_ssa_lowers_to_tfillpad.pto | 2 +- ...xpand_tile_op_tilelang_tfillpad_expand.pto | 15 ++- ...pand_tile_op_tilelang_tfillpad_inplace.pto | 3 +- .../scripts/run_remote_npu_validation.sh | 2 +- test/samples/Fillpad/fillpad_expand.py | 3 +- .../samples/Fillpad/fillpad_expand_invalid.py | 3 +- .../fillpad_expand_pad_null_invalid.py | 3 +- test/samples/Fillpad/fillpad_inplace.py | 3 +- test/samples/runop.sh | 22 +--- .../tfillpad_expand/tfillpad_expand.pto | 12 +- .../tfillpad_inplace/tfillpad_inplace.pto | 12 +- .../tfillpad_expand/tfillpad_expand.pto | 14 +- .../tfillpad_inplace/tfillpad_inplace.pto | 9 +- tools/ptobc/generated/ptobc_opcodes_v0.h | 6 - 43 files changed, 327 insertions(+), 512 deletions(-) delete mode 100644 lib/TileOps/a5/tfillpad_expand.py delete mode 100644 lib/TileOps/a5/tfillpad_inplace.py create mode 100644 test/lit/pto/tfillpad_non_normal_mat_invalid.pto diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index af326644c9..5cc540930a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -52,7 +52,7 @@ on: description: "pto-isa ref (commit/tag/branch; empty = repo-pinned weekly commit)" type: string # NOTE: Pin a known-good GitCode commit for deterministic runs. - default: ce3262e3825a235f951917eeada30e52910b6a84 + default: 27386d906e8fdcbd93aec84197939bc0b2c6caea remote_host: description: "SSH host/IP for the NPU machine" type: string @@ -427,7 +427,7 @@ jobs: SKIP_CASES: ${{ github.event.inputs.skip_cases || '' }} RUN_ONLY_CASES: ${{ github.event.inputs.run_only_cases || '' }} PTO_ISA_REPO: ${{ github.event.inputs.pto_isa_repo || 'https://gitcode.com/cann/pto-isa.git' }} - PTO_ISA_COMMIT: ${{ github.event.inputs.pto_isa_commit || 'ce3262e3825a235f951917eeada30e52910b6a84' }} + PTO_ISA_COMMIT: ${{ github.event.inputs.pto_isa_commit || '27386d906e8fdcbd93aec84197939bc0b2c6caea' }} REMOTE_HOST: ${{ github.event.inputs.remote_host || '101.245.68.6' }} REMOTE_USER: ${{ github.event.inputs.remote_user || 'zhongxuan' }} REMOTE_PORT: ${{ github.event.inputs.remote_port || '22' }} diff --git a/docker/Dockerfile b/docker/Dockerfile index 59b92eb174..045a7c3a53 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -119,8 +119,8 @@ RUN ptoas --enable-insert-sync ./abs.pto -o ./abs.cpp # NOTE: if using the smaller `ascend/python` runtime image, can only run until here, cannot test bisheng usage below # to compile cpp file generated by ptoas, need pto-isa headers -ARG PTO_ISA_COMMIT=ce3262e3825a235f951917eeada30e52910b6a84 -# pinned: https://gitcode.com/cann/pto-isa/commit/ce3262e3825a235f951917eeada30e52910b6a84 +ARG PTO_ISA_COMMIT=27386d906e8fdcbd93aec84197939bc0b2c6caea +# pinned: https://gitcode.com/cann/pto-isa/commit/27386d906e8fdcbd93aec84197939bc0b2c6caea WORKDIR /sources RUN git clone https://gitcode.com/cann/pto-isa.git \ diff --git a/docs/PTO_IR_manual.md b/docs/PTO_IR_manual.md index cd034ff079..2fb7e36efa 100644 --- a/docs/PTO_IR_manual.md +++ b/docs/PTO_IR_manual.md @@ -8038,13 +8038,14 @@ pto.textract ins(%src[%row, %col] : !pto.tile_buf<...>) outs(%dst : !pto.tile_bu ##### `pto.tfillpad` - Fill Padding Region -**Summary:** Copies `src` into `dst` and fills padded elements using `dst`'s PadVal. +**Summary:** Unified normal, in-place, and expand padding operation. `mode` defaults to `normal` and is never inferred from SSA aliasing or shapes. **Semantics:** ``` -For valid elements: dst = src -For padded elements: dst = PadVal(dst) +normal: copy valid src elements, then fill dst padding +in_place: keep valid data already in shared storage, then fill dst padding +expand: copy src into a possibly larger dst, then fill the expanded region ``` **Arguments:** @@ -8053,6 +8054,7 @@ For padded elements: dst = PadVal(dst) |------|------|-------------| | `src` | `pto.tile_buf` | Source tile | | `dst` | `pto.tile_buf` | Destination tile (with pad config) | +| `mode` | `#pto.tfillpad_mode` | PTO-ISA execution mode; defaults to `normal` | | `padValue` | `#pto.pad_value<...>` (optional) | Explicit `TFILLPAD` template argument for `loc=mat`. When present, it must match `dst`'s tile pad configuration. | **Results:** None. Writes into `dst` via DPS pattern. @@ -8061,116 +8063,30 @@ For padded elements: dst = PadVal(dst) - `dst.pad` must not be `null`. - `src` and `dst` element sizes must match, and the element size must be `1`, `2`, or `4` bytes. -- `dst.rows/cols` must match `src.rows/cols`. -- If `padValue` is present, `dst` must be `loc=mat` and `padValue` must equal the tile type's `pad`. +- Normal and in-place modes require equal source and destination static shapes. +- Expand mode requires `dst.rows >= src.rows` and `dst.cols >= src.cols`. +- Non-normal modes require both operands to use `loc=vec`. +- If `padValue` is present, mode must be normal, `dst` must be `loc=mat`, and `padValue` must equal the tile type's `pad`. - For `loc=mat`, `src` and `dst` must be lowerable to the same `TFILLPAD` tile specialization, i.e. `validShape` and `pad` must be identical. **Hardware Mapping:** -- Executes on the **Vector pipeline** (`PIPE_V`) +- VEC forms execute on the **Vector pipeline** (`PIPE_V`). +- The normal homogeneous MAT form executes on `PIPE_MTE1`. +- Normal lowers to `TFILLPAD(dst, src)`; non-normal modes lower one-to-one to `TFILLPAD(dst, src)`. **Basic Example:** ```mlir pto.tfillpad ins(%src : !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) -pto.tfillpad ins(%src : !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) - {padValue = #pto.pad_value} -``` - ---- - -##### `pto.tfillpad_expand` - Fill Padding Region With Expand - -**Summary:** Copies `src` into `dst` and fills padded elements using `dst`'s PadVal, allowing `dst` to be larger than `src`. - -**Semantics:** - -``` -For valid elements: dst = src -For padded elements: dst = PadVal(dst) -Constraint: dst.rows >= src.rows and dst.cols >= src.cols -``` - -**Arguments:** - -| Name | Type | Description | -|------|------|-------------| -| `src` | `pto.tile_buf` | Source tile | -| `dst` | `pto.tile_buf` | Destination tile (with pad config, may be larger) | - -**Results:** None. Writes into `dst` via DPS pattern. - -**Constraints & Verification:** - -- The operation has a custom verifier. -- For `loc=mat`, cross-layer behavior with heterogeneous (`src`/`dst`) expand shape is not finalized in this release; `tfillpad_expand` is not covered by the `tfillpad`-specific lowerability check. - -**Hardware Mapping:** - -- Executes on the **Vector pipeline** (`PIPE_V`) - -**Basic Example:** - -```mlir -pto.tfillpad_expand ins(%src : !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) -``` - ---- +pto.tfillpad ins(%tile : !pto.tile_buf) + outs(%tile : !pto.tile_buf) + {mode = #pto.tfillpad_mode} -##### `pto.tfillpad_inplace` - Fill Padding Region In Place - -**Summary:** Fills the padding region in place on shared backing storage. `src` provides the valid-region bounds and `dst` provides the target pad bounds/configuration. - -**Semantics:** - -``` -For elements inside src valid_shape: - dst keeps the existing value -For padded elements described by dst: - dst = PadVal(dst) -``` - -This operation is intended for the in-place case where `src` and `dst` refer to the same tile storage, often the same SSA value. - -**Arguments:** - -| Name | Type | Description | -|------|------|-------------| -| `src` | `pto.tile_buf` | Source tile supplying valid-region bounds | -| `dst` | `pto.tile_buf` | Destination tile supplying pad configuration and receiving the in-place update | - -**Results:** None. Writes into `dst` via DPS pattern. - -**Assembly Format:** - -``` -pto.tfillpad_inplace ins( : ) - outs( : ) -``` - -**Constraints & Verification:** - -- `dst.pad` must not be `null`. -- `src` and `dst` element sizes must match, and the element size must be `1`, `2`, or `4` bytes. -- `src.rows/cols` and `dst.rows/cols` must have the same static shape. -- The verifier uses the same non-expand shape constraints as `pto.tfillpad`. -- Unlike `pto.tfillpad_expand`, `dst` is not allowed to have a larger static shape than `src`. - -**Hardware Mapping:** - -- Executes on the **Vector pipeline** (`PIPE_V`) -- EmitC lowers to `TFILLPAD_INPLACE(dst, src)` - -**Basic Example:** - -```mlir -pto.tfillpad_inplace ins(%tile : !pto.tile_buf) - outs(%tile : !pto.tile_buf) +pto.tfillpad ins(%src_small : !pto.tile_buf) + outs(%dst_large : !pto.tile_buf) + {mode = #pto.tfillpad_mode} ``` --- diff --git a/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md b/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md index c089b62c5e..4750cae6fb 100644 --- a/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md +++ b/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md @@ -452,7 +452,7 @@ if opPolicy.notInplaceSafe: pto.ttrans pto.tgather pto.tands / pto.tors / pto.txors -pto.tfillpad_expand +pto.tfillpad {mode = #pto.tfillpad_mode} pto.tfmod / pto.tfmods pto.trecip / pto.trsqrt pto.trowmax / pto.trowmin / pto.trowsum / pto.trowprod @@ -461,7 +461,7 @@ pto.tcolargmax / pto.tcolargmin pto.tsort32 / pto.tmrgsort ``` -其中 `pto.tands` / `pto.tors` / `pto.txors` 和 `pto.tfillpad_expand` 是 PTOAS 侧额外保守标记的 non-inplace-safe op。它们虽然不是 scratch-output conflict，但后端/ISA 语义没有明确承诺 input/output alias 安全，memplan 不应通过地址复用隐式把它们变成 inplace 执行。 +其中 `pto.tands` / `pto.tors` / `pto.txors` 和 expand 模式的 `pto.tfillpad` 是 PTOAS 侧额外保守标记的 non-inplace-safe op。它们虽然不是 scratch-output conflict，但后端/ISA 语义没有明确承诺 input/output alias 安全，memplan 不应通过地址复用隐式把它们变成 inplace 执行。 **适用场景 sample：算法本身不支持 input/output alias。** diff --git a/docs/designs/ptoas-tile-native-mainline-op-migration.md b/docs/designs/ptoas-tile-native-mainline-op-migration.md index b66a7f5c68..f76d019dff 100644 --- a/docs/designs/ptoas-tile-native-mainline-op-migration.md +++ b/docs/designs/ptoas-tile-native-mainline-op-migration.md @@ -190,8 +190,7 @@ PTO tile/view IR | `pto.textract_fp` | fp tile role和地址空间 | | `pto.tinsert` | materialize pass 当前对 tile config 有特殊推断；目标是由 result type 完整携带 | | `pto.tinsert_fp` | fp/pre-quant tile role | -| `pto.tfillpad` | src/dst alias 和 A5 MAT/PIPE 选择 | -| `pto.tfillpad_inplace` | same-SSA inplace 和 MemoryEffects | +| `pto.tfillpad` | 显式 mode、src/dst alias、MemoryEffects 和 A5 MAT/PIPE 选择 | | `pto.tsetval` | tile writer 和 result type | | `pto.tgetval` | tile reader和 scalar result | | `pto.tgather` | optional tmp、compare/index form 和 sync macro model | diff --git a/docs/isa/tile-op/12-fill-and-padding-ops.md b/docs/isa/tile-op/12-fill-and-padding-ops.md index 0734b2f552..0900d2e179 100644 --- a/docs/isa/tile-op/12-fill-and-padding-ops.md +++ b/docs/isa/tile-op/12-fill-and-padding-ops.md @@ -3,7 +3,7 @@ > **Category:** Tile-local fill, pad, and expansion materialization > **Pipeline:** PIPE_V -This chapter documents the TileLib fill / padding families. These ops preserve or materialize valid data and then synthesize the remaining destination region from the destination tile's padding policy. +This chapter documents the unified TileLib fill / padding operation. It preserves or materializes valid data and then synthesizes the remaining destination region from the destination tile's padding policy. The destination tile's `pad` / `pad_value` configuration determines which value is written into the synthesized padding or expansion region. @@ -15,8 +15,9 @@ The destination tile's `pad` / `pad_value` configuration determines which value ```mlir pto.tfillpad ins(%src : !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) + {mode = #pto.tfillpad_mode} ``` -- **semantics:** copy valid data from `src` into `dst`, then fill the remaining destination region according to `dst`'s pad policy. +- **semantics:** the `mode` attribute selects normal, in-place, or expand behavior. It defaults to `normal`; PTOAS does not infer it from aliasing or shape. **Parameter Table:** @@ -24,78 +25,36 @@ pto.tfillpad ins(%src : !pto.tile_buf<...>) |-----------|------|-------------| | `src` | `pto.tile_buf` | Source tile. | | `dst` | `pto.tile_buf` | Destination tile carrying the pad configuration. | +| `mode` | `#pto.tfillpad_mode` | ISA mode; defaults to `normal`. | +| `padValue` | `#pto.pad_value<...>` (optional) | Explicit MAT `TFILLPAD` argument; only valid in normal mode. | + +**Mode Table:** + +| Mode | Behavior | PTO-ISA mapping | +|------|----------|-----------------| +| `normal` | Copy valid data from `src`, then fill padding in `dst`. | `TFILLPAD(dst, src)` | +| `in_place` | Skip the copy phase and fill padding on shared storage. | `TFILLPAD(dst, src)` | +| `expand` | Copy `src` into a destination whose static shape may be larger, then fill the expanded region. | `TFILLPAD(dst, src)` | **Constraints:** - Source and destination element types must be compatible. - The destination tile must carry a meaningful pad configuration. -- This family is VEC-oriented. +- `in_place` and `expand` are VEC-only. Normal mode also supports the homogeneous MAT overload. +- Normal and in-place modes require equal source and destination static shapes. +- Expand mode requires each destination static dimension to be greater than or equal to the source dimension. **Example:** ```mlir pto.tfillpad ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) -``` - ---- - -## 12.2 `pto.tfillpad_expand` -- **syntax:** -```mlir -pto.tfillpad_expand ins(%src : !pto.tile_buf<...>) - outs(%dst : !pto.tile_buf<...>) -``` -- **semantics:** copy valid data from `src` into `dst`, then fill row/column expansion according to `dst`'s pad policy when the destination valid region or backing shape is larger than the source. - -**Parameter Table:** +pto.tfillpad ins(%tile : !pto.tile_buf) + outs(%tile : !pto.tile_buf) + {mode = #pto.tfillpad_mode} -| Parameter | Type | Description | -|-----------|------|-------------| -| `src` | `pto.tile_buf` | Source tile. | -| `dst` | `pto.tile_buf` | Larger destination tile carrying the pad configuration. | - -**Constraints:** - -- `dst` may be larger than `src` in valid region or physical shape. -- The fill value is derived from `dst.pad_value`. -- A unified VEC-oriented template handles the supported element families. - -**Example:** - -```mlir -pto.tfillpad_expand ins(%src : !pto.tile_buf) - outs(%dst : !pto.tile_buf) -``` - ---- - -## 12.3 `pto.tfillpad_inplace` - -- **syntax:** -```mlir -pto.tfillpad_inplace ins(%src : !pto.tile_buf<...>) - outs(%dst : !pto.tile_buf<...>) -``` -- **semantics:** update the padding / expansion region of an already materialized tile without a separate copy-in phase. - -**Parameter Table:** - -| Parameter | Type | Description | -|-----------|------|-------------| -| `src` | `pto.tile_buf` | Source tile buffer. | -| `dst` | `pto.tile_buf` | Destination tile buffer, typically aliasing the same physical tile. | - -**Constraints:** - -- PTOAS exposes `pto.tfillpad_inplace` as a dedicated Tile op. -- In typical use, `src` and `dst` refer to the same underlying tile buffer. -- The fill value is derived from `dst.pad_value`. - -**Example:** - -```mlir -pto.tfillpad_inplace ins(%tile : !pto.tile_buf) - outs(%tile : !pto.tile_buf) +pto.tfillpad ins(%src_small : !pto.tile_buf) + outs(%dst_large : !pto.tile_buf) + {mode = #pto.tfillpad_mode} ``` diff --git a/docs/release/PTO-tile-Instruction-SPEC-v0.4.md b/docs/release/PTO-tile-Instruction-SPEC-v0.4.md index 2ed51a07ca..a3448620dc 100644 --- a/docs/release/PTO-tile-Instruction-SPEC-v0.4.md +++ b/docs/release/PTO-tile-Instruction-SPEC-v0.4.md @@ -1757,7 +1757,7 @@ pto.tsels ins(%mask, %src, %tmp, %scalar : > **Category:** Tile-local fill, pad, and expansion materialization > **Pipeline:** PIPE_V -This chapter documents the TileLib fill / padding families. These ops preserve or materialize valid data and then synthesize the remaining destination region from the destination tile's padding policy. +This chapter documents the unified TileLib fill / padding operation. It preserves or materializes valid data and then synthesizes the remaining destination region from the destination tile's padding policy. The destination tile's `pad` / `pad_value` configuration determines which value is written into the synthesized padding or expansion region. @@ -1769,8 +1769,9 @@ The destination tile's `pad` / `pad_value` configuration determines which value ```mlir pto.tfillpad ins(%src : !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) + {mode = #pto.tfillpad_mode} ``` -- **semantics:** copy valid data from `src` into `dst`, then fill the remaining destination region according to `dst`'s pad policy. +- **semantics:** `mode` selects normal, in-place, or expand behavior and defaults to `normal`. PTOAS does not infer the mode from shapes or aliasing. **Parameter Table:** @@ -1778,47 +1779,34 @@ pto.tfillpad ins(%src : !pto.tile_buf<...>) |-----------|------|-------------| | `src` | `pto.tile_buf` | Source tile. | | `dst` | `pto.tile_buf` | Destination tile carrying the pad configuration. | +| `mode` | `#pto.tfillpad_mode` | PTO-ISA execution mode. | + +**Mode Table:** + +| Mode | Behavior | PTO-ISA mapping | +|------|----------|-----------------| +| `normal` | Copy valid data, then fill padding. | `TFILLPAD(dst, src)` | +| `in_place` | Skip the copy phase and fill padding on shared storage. | `TFILLPAD(dst, src)` | +| `expand` | Copy into a possibly larger destination and fill the expanded region. | `TFILLPAD(dst, src)` | **Constraints:** - Source and destination element types must be compatible. - The destination tile must carry a meaningful pad configuration. -- This family is VEC-oriented. +- Non-normal modes are VEC-only. Normal mode also supports the homogeneous MAT overload. +- Normal and in-place modes require equal static shapes; expand requires each destination dimension to be greater than or equal to the source dimension. **Example:** ```mlir pto.tfillpad ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) -``` - ---- - -### 12.2 `pto.tfillpad_expand` -- **syntax:** -```mlir -pto.tfillpad_expand ins(%src : !pto.tile_buf<...>) - outs(%dst : !pto.tile_buf<...>) -``` -- **semantics:** copy valid data from `src` into `dst`, then fill row/column expansion according to `dst`'s pad policy when the destination valid region or backing shape is larger than the source. - -**Parameter Table:** +pto.tfillpad ins(%tile : !pto.tile_buf) + outs(%tile : !pto.tile_buf) + {mode = #pto.tfillpad_mode} -| Parameter | Type | Description | -|-----------|------|-------------| -| `src` | `pto.tile_buf` | Source tile. | -| `dst` | `pto.tile_buf` | Larger destination tile carrying the pad configuration. | - -**Constraints:** - -- `dst` may be larger than `src` in valid region or physical shape. -- The fill value is derived from `dst.pad_value`. -- A unified VEC-oriented template handles the supported element families. - -**Example:** - -```mlir -pto.tfillpad_expand ins(%src : !pto.tile_buf) - outs(%dst : !pto.tile_buf) +pto.tfillpad ins(%src_small : !pto.tile_buf) + outs(%dst_large : !pto.tile_buf) + {mode = #pto.tfillpad_mode} ``` diff --git a/include/PTO/IR/PTOAttrs.td b/include/PTO/IR/PTOAttrs.td index d597173105..fdf95e0cea 100644 --- a/include/PTO/IR/PTOAttrs.td +++ b/include/PTO/IR/PTOAttrs.td @@ -1190,6 +1190,19 @@ def PTO_PadValueAttr : PTO_Attr<"PadValue", "pad_value"> { let assemblyFormat = "`<` params `>`"; } +def PTO_TFillPadModeEnum : PTO_I32Enum< + "TFillPadMode", "PTO TFILLPAD execution mode", [ + I32EnumAttrCase<"Normal", 0, "normal">, + I32EnumAttrCase<"InPlace", 1, "in_place">, + I32EnumAttrCase<"Expand", 2, "expand"> + ]>; + +def PTO_TFillPadModeAttr + : EnumAttr { + let assemblyFormat = "`<` params `>`"; + let summary = "TFILLPAD normal, in-place, or expand execution mode"; +} + def PTO_CompactMode_Enum : PTO_I32Enum<"CompactMode", "Tile compact mode", [ I32EnumAttrCase<"Null", 0, "null">, I32EnumAttrCase<"Normal", 1, "normal">, diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 30bcc3d4b6..5c29a5f9a6 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -4750,12 +4750,13 @@ def TFillPadOp : PTO_TOp<"tfillpad", [ OpPipeInterface, DeclareOpInterfaceMethods ]> { - let summary = "Copy src into dst and fill padded elements using dst PadVal (tilebuf, DPS)"; + let summary = "Fill padding in normal, in-place, or expand mode (tilebuf, DPS)"; let arguments = (ins PTODpsType:$src, PTODpsType:$dst, - OptionalAttr:$padValue + OptionalAttr:$padValue, + DefaultValuedAttr:$mode ); let results = (outs); @@ -4793,62 +4794,6 @@ def TFillPadOp : PTO_TOp<"tfillpad", [ }]; } -def TFillPadExpandOp : PTO_TOp<"tfillpad_expand", [ - PTO_DpsInitOpInterface, - OpPipeInterface, - DeclareOpInterfaceMethods -]> { - let summary = "Copy src into dst and fill padded elements using dst PadVal, allowing dst to be larger (tilebuf, DPS)"; - - let arguments = (ins - PTODpsType:$src, - PTODpsType:$dst - ); - - let results = (outs); - - let hasVerifier = 1; - - let assemblyFormat = [{ - `ins` `(` $src `:` qualified(type($src)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; - - let extraClassDeclaration = [{ - ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } - ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } - }]; -} - -def TFillPadInplaceOp : PTO_TOp<"tfillpad_inplace", [ - PTO_DpsInitOpInterface, - OpPipeInterface, - DeclareOpInterfaceMethods -]> { - let summary = "In-place fill padding on shared backing storage: src supplies valid-region bounds, dst supplies target pad bounds (tilebuf, DPS)"; - - let arguments = (ins - PTODpsType:$src, - PTODpsType:$dst - ); - - let results = (outs); - - let hasVerifier = 1; - - let assemblyFormat = [{ - `ins` `(` $src `:` qualified(type($src)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; - - let extraClassDeclaration = [{ - ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } - ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } - }]; -} - def TGatherOp : PTO_TOp<"tgather", [ AttrSizedOperandSegments, PTO_DpsInitOpInterface, diff --git a/include/pto-c/Dialect/PTO.h b/include/pto-c/Dialect/PTO.h index cb653c2384..e81917467f 100644 --- a/include/pto-c/Dialect/PTO.h +++ b/include/pto-c/Dialect/PTO.h @@ -109,6 +109,9 @@ MLIR_CAPI_EXPORTED int32_t mlirPTOPadValueAttrGetValue(MlirAttribute attr); MLIR_CAPI_EXPORTED bool mlirPTOAttrIsACompactModeAttr(MlirAttribute attr); MLIR_CAPI_EXPORTED MlirAttribute mlirPTOCompactModeAttrGet(MlirContext ctx, int32_t value); MLIR_CAPI_EXPORTED int32_t mlirPTOCompactModeAttrGetValue(MlirAttribute attr); +MLIR_CAPI_EXPORTED bool mlirPTOAttrIsATFillPadModeAttr(MlirAttribute attr); +MLIR_CAPI_EXPORTED MlirAttribute mlirPTOTFillPadModeAttrGet(MlirContext ctx, int32_t value); +MLIR_CAPI_EXPORTED int32_t mlirPTOTFillPadModeAttrGetValue(MlirAttribute attr); MLIR_CAPI_EXPORTED bool mlirPTOAttrIsAAccToVecModeAttr(MlirAttribute attr); MLIR_CAPI_EXPORTED MlirAttribute mlirPTOAccToVecModeAttrGet(MlirContext ctx, int32_t value); MLIR_CAPI_EXPORTED int32_t mlirPTOAccToVecModeAttrGetValue(MlirAttribute attr); diff --git a/lib/Bindings/Python/PTOModule.cpp b/lib/Bindings/Python/PTOModule.cpp index b8377ca585..4d75e1d36d 100644 --- a/lib/Bindings/Python/PTOModule.cpp +++ b/lib/Bindings/Python/PTOModule.cpp @@ -257,6 +257,12 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { .value("DualModeSplitN", mlir::pto::AccToVecMode::DualModeSplitN) .export_values(); + py::enum_(m, "TFillPadMode") + .value("Normal", mlir::pto::TFillPadMode::Normal) + .value("InPlace", mlir::pto::TFillPadMode::InPlace) + .value("Expand", mlir::pto::TFillPadMode::Expand) + .export_values(); + py::enum_(m, "TInsertMode") .value("SPLIT2", mlir::pto::TInsertMode::SPLIT2) .value("SPLIT4", mlir::pto::TInsertMode::SPLIT4) @@ -393,6 +399,19 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { }, py::arg("cls"), py::arg("value"), py::arg("context") = py::none()); + mlir_attribute_subclass(m, "TFillPadModeAttr", + [](MlirAttribute a) -> bool { + return mlirPTOAttrIsATFillPadModeAttr(a); + }) + .def_classmethod( + "get", + [](py::object cls, mlir::pto::TFillPadMode value, MlirContext ctx) -> py::object { + MlirAttribute a = mlirPTOTFillPadModeAttrGet(ctx, static_cast(value)); + if (mlirAttributeIsNull(a)) return py::none(); + return cls(a); + }, + py::arg("cls"), py::arg("value"), py::arg("context") = py::none()); + mlir_attribute_subclass(m, "TInsertModeAttr", [](MlirAttribute a) -> bool { return mlirPTOAttrIsATInsertModeAttr(a); diff --git a/lib/CAPI/Dialect/PTO.cpp b/lib/CAPI/Dialect/PTO.cpp index 40520b90fc..9bea9f9433 100644 --- a/lib/CAPI/Dialect/PTO.cpp +++ b/lib/CAPI/Dialect/PTO.cpp @@ -734,6 +734,21 @@ int32_t mlirPTOCompactModeAttrGetValue(MlirAttribute attr) { return static_cast(a.getValue()); } +bool mlirPTOAttrIsATFillPadModeAttr(MlirAttribute attr) { + return mlir::isa(unwrap(attr)); +} + +MlirAttribute mlirPTOTFillPadModeAttrGet(MlirContext ctx, int32_t value) { + auto *c = unwrap(ctx); + return wrap(mlir::pto::TFillPadModeAttr::get( + c, static_cast(value))); +} + +int32_t mlirPTOTFillPadModeAttrGetValue(MlirAttribute attr) { + auto a = mlir::cast(unwrap(attr)); + return static_cast(a.getValue()); +} + bool mlirPTOAttrIsAAccToVecModeAttr(MlirAttribute attr) { return mlir::isa(unwrap(attr)); } diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index aa74bfb955..bf6caf9b4e 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -7269,9 +7269,9 @@ mlir::LogicalResult mlir::pto::TInsertFPOp::verify() { return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); } -static mlir::LogicalResult verifyTFillPadLike(Operation *op, Type srcTy, Type dstTy, - bool allowDstExpand, - llvm::StringRef opName) { +static mlir::LogicalResult verifyTFillPadLike(Operation *op, Type srcTy, + Type dstTy, + pto::TFillPadMode mode) { if (!isPTOShapedLike(srcTy) || !isPTOShapedLike(dstTy)) return op->emitError("expects src/dst to be PTO shaped-like types"); @@ -7297,67 +7297,68 @@ static mlir::LogicalResult verifyTFillPadLike(Operation *op, Type srcTy, Type ds if (!(srcB == 1 || srcB == 2 || srcB == 4)) return op->emitError("expects element size to be 1, 2, or 4 bytes"); + auto srcSpace = getPTOMemorySpaceEnum(srcTy); + auto dstSpace = getPTOMemorySpaceEnum(dstTy); + if (mode != pto::TFillPadMode::Normal && + (!srcSpace || !dstSpace || *srcSpace != pto::AddressSpace::VEC || + *dstSpace != pto::AddressSpace::VEC)) + return op->emitError() + << "expects non-normal TFILLPAD mode only for loc=vec"; + // pto.tfillpad lowers to TFILLPAD(dst, src). For loc=mat, pto-isa only // exposes the homogeneous overload, so src/dst must use the same Tile<...> // specialization (including valid_shape and pad). - // Note: tfillpad_expand is intentionally not covered here because its - // cross-layer ABI contract for loc=mat heterogeneous shape expansion is not - // finalized yet. - if (opName == "tfillpad") { - auto srcTb = mlir::dyn_cast(srcTy); - auto dstTb = mlir::dyn_cast(dstTy); - auto srcSpace = getPTOMemorySpaceEnum(srcTy); - auto dstSpace = getPTOMemorySpaceEnum(dstTy); - if (srcTb && dstTb && srcSpace && dstSpace && - *srcSpace == mlir::pto::AddressSpace::MAT && - *dstSpace == mlir::pto::AddressSpace::MAT && srcTb != dstTb) { - auto dimToStr = [](int64_t dim) -> std::string { - return dim == ShapedType::kDynamic ? "?" : std::to_string(dim); - }; - SmallVector mismatchFields; - auto srcValid = getValidShapeVec(srcTy); - auto dstValid = getValidShapeVec(dstTy); - if (srcValid.size() == 2 && dstValid.size() == 2) { - if (srcValid[0] != dstValid[0]) - mismatchFields.push_back("v_row (" + dimToStr(srcValid[0]) + " vs " + - dimToStr(dstValid[0]) + ")"); - if (srcValid[1] != dstValid[1]) - mismatchFields.push_back("v_col (" + dimToStr(srcValid[1]) + " vs " + - dimToStr(dstValid[1]) + ")"); - } - if (srcTb.getPadValueI32() != dstTb.getPadValueI32()) - mismatchFields.push_back("pad (" + std::to_string(srcTb.getPadValueI32()) + - " vs " + std::to_string(dstTb.getPadValueI32()) + - ")"); - - auto diag = op->emitError() - << "expects src/dst tile types to be lowerable to TFILLPAD " - "for loc=mat"; - if (!mismatchFields.empty()) - diag << "; mismatching fields: " << llvm::join(mismatchFields, ", "); - diag << "\n src: " << srcTy; - diag << "\n dst: " << dstTy; - diag << "\n note: heterogeneous TFILLPAD overload is only available for loc=vec"; - return failure(); + auto srcTb = mlir::dyn_cast(srcTy); + auto dstTb = mlir::dyn_cast(dstTy); + if (srcTb && dstTb && srcSpace && dstSpace && + *srcSpace == mlir::pto::AddressSpace::MAT && + *dstSpace == mlir::pto::AddressSpace::MAT && srcTb != dstTb) { + auto dimToStr = [](int64_t dim) -> std::string { + return dim == ShapedType::kDynamic ? "?" : std::to_string(dim); + }; + SmallVector mismatchFields; + auto srcValid = getValidShapeVec(srcTy); + auto dstValid = getValidShapeVec(dstTy); + if (srcValid.size() == 2 && dstValid.size() == 2) { + if (srcValid[0] != dstValid[0]) + mismatchFields.push_back("v_row (" + dimToStr(srcValid[0]) + " vs " + + dimToStr(dstValid[0]) + ")"); + if (srcValid[1] != dstValid[1]) + mismatchFields.push_back("v_col (" + dimToStr(srcValid[1]) + " vs " + + dimToStr(dstValid[1]) + ")"); } + if (srcTb.getPadValueI32() != dstTb.getPadValueI32()) + mismatchFields.push_back("pad (" + std::to_string(srcTb.getPadValueI32()) + + " vs " + std::to_string(dstTb.getPadValueI32()) + + ")"); + + auto diag = op->emitError() + << "expects src/dst tile types to be lowerable to TFILLPAD " + "for loc=mat"; + if (!mismatchFields.empty()) + diag << "; mismatching fields: " << llvm::join(mismatchFields, ", "); + diag << "\n src: " << srcTy; + diag << "\n dst: " << dstTy; + diag << "\n note: heterogeneous TFILLPAD overload is only available for loc=vec"; + return failure(); } if (auto dstTileTy = mlir::dyn_cast(dstTy)) { auto padAttr = mlir::dyn_cast(dstTileTy.getPadValueAttr()); if (!padAttr || padAttr.getValue() == mlir::pto::PadValue::Null) - return op->emitError() << "expects dst PadVal != Null for " << opName; + return op->emitError("expects dst PadVal != Null for tfillpad"); } - if (!allowDstExpand) { + if (mode != pto::TFillPadMode::Expand) { if (srcShape != dstShape) - return op->emitError() - << "expects src and dst to have the same static shape for " << opName; + return op->emitError("expects src and dst to have the same static shape " + "unless mode is expand"); return mlir::success(); } if (srcShape[0] > dstShape[0] || srcShape[1] > dstShape[1]) { - return op->emitError() - << "expects dst static shape to be >= src static shape for " << opName; + return op->emitError( + "expects dst static shape to be >= src static shape for expand mode"); } return mlir::success(); @@ -7365,10 +7366,12 @@ static mlir::LogicalResult verifyTFillPadLike(Operation *op, Type srcTy, Type ds mlir::LogicalResult mlir::pto::TFillPadOp::verify() { if (failed(verifyTFillPadLike(getOperation(), getSrc().getType(), getDst().getType(), - /*allowDstExpand=*/false, "tfillpad"))) + getMode()))) return failure(); if (auto padValueAttr = getPadValueAttr()) { + if (getMode() != pto::TFillPadMode::Normal) + return emitOpError("expects padValue attribute only for normal mode"); auto dstSpace = getPTOMemorySpaceEnum(getDst().getType()); if (!dstSpace || *dstSpace != pto::AddressSpace::MAT) return emitOpError("expects padValue attribute only for loc=mat tfillpad"); @@ -7382,16 +7385,6 @@ mlir::LogicalResult mlir::pto::TFillPadOp::verify() { return success(); } -mlir::LogicalResult mlir::pto::TFillPadExpandOp::verify() { - return verifyTFillPadLike(getOperation(), getSrc().getType(), getDst().getType(), - /*allowDstExpand=*/true, "tfillpad_expand"); -} - -mlir::LogicalResult mlir::pto::TFillPadInplaceOp::verify() { - return verifyTFillPadLike(getOperation(), getSrc().getType(), getDst().getType(), - /*allowDstExpand=*/false, "tfillpad_inplace"); -} - llvm::LogicalResult mlir::pto::TGatherOp::verify() { auto isSupportedGatherElemTypeA5Index = [&](Type ty) -> bool { @@ -14136,8 +14129,6 @@ void TInsertFPOp::getEffects( } PTO_DEFINE_UNARY_EFFECTS(TFillPadOp, getSrcMutable(), getDstMutable()) -PTO_DEFINE_UNARY_EFFECTS(TFillPadExpandOp, getSrcMutable(), getDstMutable()) -PTO_DEFINE_UNARY_EFFECTS(TFillPadInplaceOp, getSrcMutable(), getDstMutable()) void TGatherOp::getEffects( SmallVectorImpl> &effects) { diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 750b9863ee..afc6d9cdc0 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -261,7 +261,7 @@ static InplacePolicy getInplacePolicy(Operation *op) { policy.notInplaceSafe = isOneOf( name, { - "pto.tands", "pto.tfillpad_expand", "pto.tfmod", + "pto.tands", "pto.tfmod", "pto.tfmods", "pto.tgather", "pto.tmrgsort", "pto.tors", "pto.trecip", "pto.trsqrt", "pto.tsort32", "pto.ttrans", "pto.trowargmax", @@ -270,6 +270,9 @@ static InplacePolicy getInplacePolicy(Operation *op) { "pto.tcolargmin", "pto.tcvt", "pto.txors", }); + if (auto fillPad = dyn_cast(op)) + policy.notInplaceSafe |= fillPad.getMode() == TFillPadMode::Expand; + if (name == "pto.tsel") { policy.forbidOutputAliasOperands.push_back(0); // mask policy.forbidOutputAliasOperands.push_back(3); // tmp diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 6b93f1d60c..2d5bd758e1 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -9824,6 +9824,18 @@ struct PTOInsertFPToEmitC : public OpConversionPattern { // pto.tfillpad lowering -> TFILLPAD(dst, src) //===----------------------------------------------------------------------===// +static StringRef getTFillPadModeToken(pto::TFillPadMode mode) { + switch (mode) { + case pto::TFillPadMode::Normal: + return "pto::TFillPadMode::Normal"; + case pto::TFillPadMode::InPlace: + return "pto::TFillPadMode::InPlace"; + case pto::TFillPadMode::Expand: + return "pto::TFillPadMode::Expand"; + } + llvm_unreachable("unknown TFillPadMode"); +} + struct PTOFillPadToEmitC : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -9855,6 +9867,9 @@ struct PTOFillPadToEmitC : public OpConversionPattern { // tfillpad, so lowering can trust the preserved semantic contract. templateArgs = rewriter.getArrayAttr( {emitc::OpaqueAttr::get(ctx, padValueTok(padValueAttr.getValue()))}); + } else if (op.getMode() != pto::TFillPadMode::Normal) { + templateArgs = rewriter.getArrayAttr( + {emitc::OpaqueAttr::get(ctx, getTFillPadModeToken(op.getMode()))}); } rewriter.create( @@ -9867,54 +9882,6 @@ struct PTOFillPadToEmitC : public OpConversionPattern { } }; //===----------------------------------------------------------------------===// -// pto.tfillpad_inplace lowering -> TFILLPAD_INPLACE(dst, src) -//===----------------------------------------------------------------------===// - -struct PTOFillPadInplaceToEmitC - : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult matchAndRewrite(pto::TFillPadInplaceOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto loc = op.getLoc(); - - Value src = peelUnrealized(adaptor.getSrc()); - Value dst = peelUnrealized(adaptor.getDst()); - - rewriter.create( - loc, TypeRange{}, "TFILLPAD_INPLACE", - /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, - /*operands=*/ValueRange{dst, src}); - - rewriter.eraseOp(op); - return success(); - } -}; -//===----------------------------------------------------------------------===// -// pto.tfillpad_expand lowering -> TFILLPAD_EXPAND(dst, src) -//===----------------------------------------------------------------------===// - -struct PTOFillPadExpandToEmitC - : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult matchAndRewrite(pto::TFillPadExpandOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto loc = op.getLoc(); - - Value src = peelUnrealized(adaptor.getSrc()); - Value dst = peelUnrealized(adaptor.getDst()); - - rewriter.create( - loc, TypeRange{}, "TFILLPAD_EXPAND", - /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, - /*operands=*/ValueRange{dst, src}); - - rewriter.eraseOp(op); - return success(); - } -}; -//===----------------------------------------------------------------------===// // pto.tgather lowering // - Index form : TGATHER(dst, src0, indices, tmp) // - Compare form: TGATHER(dst, src0, kValue, cdst, tmp) @@ -13460,8 +13427,7 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); - patterns.add( - typeConverter, ctx); + patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); diff --git a/lib/TileOps/__init__.py b/lib/TileOps/__init__.py index cadc622104..f595bfd22d 100644 --- a/lib/TileOps/__init__.py +++ b/lib/TileOps/__init__.py @@ -46,8 +46,6 @@ ("a5", "pto.tfmod"): ".a5.tfmod", ("a5", "pto.tfmods"): ".a5.tfmods", ("a5", "pto.tfillpad"): ".a5.tfillpad", - ("a5", "pto.tfillpad_expand"): ".a5.tfillpad_expand", - ("a5", "pto.tfillpad_inplace"): ".a5.tfillpad_inplace", ("a5", "pto.tgatherb"): ".a5.tgatherb", ("a5", "pto.tgemv"): ".a5.tgemv", ("a5", "pto.tgemv.acc"): ".a5.tgemv_acc", diff --git a/lib/TileOps/a5/_fillpad.py b/lib/TileOps/a5/_fillpad.py index 9c7198654d..5e56896e47 100644 --- a/lib/TileOps/a5/_fillpad.py +++ b/lib/TileOps/a5/_fillpad.py @@ -152,11 +152,11 @@ def _fill_inplace(dst, src_valid_rows, src_valid_cols, dst_valid_rows, dst_valid _fill(dst, src_valid_rows, dst_valid_rows, 0, dst_valid_cols) -def register_fillpad(*, op, name, copy): +def register_fillpad(): @tilelib.tile_template( - op=op, + op="pto.tfillpad", target="a5", - name=name, + name="template_tfillpad", dtypes=_DTYPES, iteration_axis="none", op_engine="other", @@ -168,20 +168,19 @@ def register_fillpad(*, op, name, copy): tags=("fillpad",), ) def template(src: pto.Tile, dst: pto.Tile): + mode = pto.get_op_attr("mode", "normal") src_valid_rows, src_valid_cols = src.valid_shape dst_valid_rows, dst_valid_cols = dst.valid_shape lanes = pto.elements_per_vreg(dst.dtype) aligned_cols = (src_valid_cols // lanes) * lanes - if not copy: + if mode == "in_place": _fill_inplace(dst, src_valid_rows, src_valid_cols, dst_valid_rows, dst_valid_cols) return - if copy: - _copy_region(src, dst, src_valid_rows, 0, aligned_cols) - fill_row_stop = dst_valid_rows if op == "pto.tfillpad_expand" else src_valid_rows + _copy_region(src, dst, src_valid_rows, 0, aligned_cols) + fill_row_stop = dst_valid_rows if mode == "expand" else src_valid_rows scalar_tail_start = _scalar_tail_start(dst, lanes) _fill(dst, 0, fill_row_stop, aligned_cols, dst_valid_cols, scalar_tail_start=scalar_tail_start) - if copy: - _copy_region(src, dst, src_valid_rows, aligned_cols, src_valid_cols) + _copy_region(src, dst, src_valid_rows, aligned_cols, src_valid_cols) _fill(dst, src_valid_rows, dst_valid_rows, 0, dst_valid_cols, scalar_tail_start=scalar_tail_start) return template diff --git a/lib/TileOps/a5/tfillpad.py b/lib/TileOps/a5/tfillpad.py index 97ae628b69..a8c7c20c22 100644 --- a/lib/TileOps/a5/tfillpad.py +++ b/lib/TileOps/a5/tfillpad.py @@ -10,4 +10,4 @@ from ._fillpad import register_fillpad -template_tfillpad = register_fillpad(op="pto.tfillpad", name="template_tfillpad", copy=True) +template_tfillpad = register_fillpad() diff --git a/lib/TileOps/a5/tfillpad_expand.py b/lib/TileOps/a5/tfillpad_expand.py deleted file mode 100644 index 9631096893..0000000000 --- a/lib/TileOps/a5/tfillpad_expand.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software, you can redistribute it and/or modify it under the terms and conditions of -# CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -# See LICENSE in the root of the software repository for the full text of the License. -"""PTODSL TileLib template for ``pto.tfillpad_expand``.""" - -from ._fillpad import register_fillpad - - -template_tfillpad_expand = register_fillpad( - op="pto.tfillpad_expand", - name="template_tfillpad_expand", - copy=True, -) diff --git a/lib/TileOps/a5/tfillpad_inplace.py b/lib/TileOps/a5/tfillpad_inplace.py deleted file mode 100644 index 533ffc92d3..0000000000 --- a/lib/TileOps/a5/tfillpad_inplace.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software, you can redistribute it and/or modify it under the terms and conditions of -# CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -# See LICENSE in the root of the software repository for the full text of the License. -"""PTODSL TileLib template for ``pto.tfillpad_inplace``.""" - -from ._fillpad import register_fillpad - - -template_tfillpad_inplace = register_fillpad( - op="pto.tfillpad_inplace", - name="template_tfillpad_inplace", - copy=False, -) diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 07b70d9583..3e2d4593f3 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -4007,27 +4007,31 @@ def tpartmin(src0, src1, dst): ) -def tfillpad(src, dst): - """``pto.tfillpad ins(src) outs(dst)``.""" - _pto.tfillpad( - unwrap_surface_value(src), - unwrap_surface_value(dst), - ) - - -def tfillpad_expand(src, dst): - """``pto.tfillpad_expand ins(src) outs(dst)``.""" - _pto.tfillpad_expand( - unwrap_surface_value(src), - unwrap_surface_value(dst), - ) +def _tfillpad_mode_attr(mode): + if isinstance(mode, Attribute): + return mode + if isinstance(mode, str): + token = mode.strip().lower().replace("-", "_") + aliases = { + "normal": _pto.TFillPadMode.Normal, + "inplace": _pto.TFillPadMode.InPlace, + "in_place": _pto.TFillPadMode.InPlace, + "expand": _pto.TFillPadMode.Expand, + } + if token not in aliases: + raise ValueError( + "tfillpad mode must be 'normal', 'in_place', or 'expand'" + ) + mode = aliases[token] + return _pto.TFillPadModeAttr.get(mode) -def tfillpad_inplace(src, dst): - """``pto.tfillpad_inplace ins(src) outs(dst)``.""" - _pto.tfillpad_inplace( +def tfillpad(src, dst, *, mode="normal"): + """``pto.tfillpad ins(src) outs(dst)`` with an explicit ISA mode.""" + _pto.tfillpad( unwrap_surface_value(src), unwrap_surface_value(dst), + mode=_tfillpad_mode_attr(mode), ) @@ -6348,7 +6352,7 @@ def import_reserved_buffer(name, *, peer_func): "tsel", "tsels", "tcvt", "tnot", "tand", "tands", "tor", "tors", "txor", "txors", "tshl", "tshls", "tshr", "tshrs", "tpartadd", "tpartmul", "tpartmax", "tpartmin", - "tfillpad", "tfillpad_expand", "tfillpad_inplace", + "tfillpad", "ttri", "tthistogram", "chistv2", "as_ptr", diff --git a/ptodsl/ptodsl/_tile_namespace.py b/ptodsl/ptodsl/_tile_namespace.py index 26cbaeafef..72995ff371 100644 --- a/ptodsl/ptodsl/_tile_namespace.py +++ b/ptodsl/ptodsl/_tile_namespace.py @@ -207,8 +207,6 @@ def rowargmin(src, dst, *, tmp=None): partmin = staticmethod(_ops.tpartmin) fillpad = staticmethod(_ops.tfillpad) - fillpad_expand = staticmethod(_ops.tfillpad_expand) - fillpad_inplace = staticmethod(_ops.tfillpad_inplace) tile = _TileNamespace() diff --git a/ptodsl/tests/test_tilelib_catalog.py b/ptodsl/tests/test_tilelib_catalog.py index 373304b0eb..24c2c1e9c1 100644 --- a/ptodsl/tests/test_tilelib_catalog.py +++ b/ptodsl/tests/test_tilelib_catalog.py @@ -75,8 +75,6 @@ "pto.tfmod": ("template_tfmod", "pto.vtrc", ("src0", "src1", "dst"), "f32"), "pto.tfmods": ("template_tfmods", "pto.vtrc", ("src", "scalar", "dst"), "f32"), "pto.tfillpad": ("template_tfillpad", "pto.vsts", ("src", "dst"), "f32"), - "pto.tfillpad_expand": ("template_tfillpad_expand", "pto.vsts", ("src", "dst"), "f32"), - "pto.tfillpad_inplace": ("template_tfillpad_inplace", "pto.vdup", ("src", "dst"), "f32"), "pto.tgemv": ("template_tgemv", "pto.mad", ("lhs", "rhs", "acc"), "f16"), "pto.tgemv.acc": ("template_tgemv_acc", "pto.mad_acc", ("acc_in", "lhs", "rhs", "dst"), "f16"), "pto.tgemv.bias": ("template_tgemv_bias", "pto.mad_bias", ("lhs", "rhs", "bias", "dst"), "f16"), @@ -269,7 +267,6 @@ ) OPS_WITHOUT_TILE_LOAD = {"pto.texpands"} OPS_WITHOUT_TILE_LOAD = OPS_WITHOUT_TILE_LOAD | {"pto.trandom", "pto.tsort32", "pto.tload", "pto.tstore", "pto.tstore_fp", "pto.textract_fp"} -OPS_WITHOUT_TILE_LOAD = OPS_WITHOUT_TILE_LOAD | {"pto.tfillpad_inplace"} OPS_WITHOUT_TILE_LOAD = OPS_WITHOUT_TILE_LOAD | CUBE_OPS OPS_WITHOUT_VECTOR_STORE = {"pto.tcmp", "pto.tcmps", "pto.tsort32"} OPS_WITHOUT_VECTOR_STORE = OPS_WITHOUT_VECTOR_STORE | {"pto.tload", "pto.tstore", "pto.tstore_fp", "pto.textract_fp"} diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index afbcb0d4c3..41dcaa0819 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -118,11 +118,12 @@ def test_tile_bitwise_aliases_are_exposed_without_legacy_names(self): def test_tile_partial_and_fillpad_names_are_exposed_without_legacy_names(self): preferred_names = [ "partadd", "partmul", "partmax", "partmin", - "fillpad", "fillpad_expand", "fillpad_inplace", + "fillpad", ] legacy_names = [ "part_add", "part_mul", "part_max", "part_min", "fill_pad", "fill_pad_expand", "fill_pad_inplace", + "fillpad_expand", "fillpad_inplace", ] for name in preferred_names: @@ -133,6 +134,23 @@ def test_tile_partial_and_fillpad_names_are_exposed_without_legacy_names(self): with self.subTest(name=name): self.assertFalse(hasattr(pto.tile, name), name) + def test_tile_fillpad_dispatches_one_op_with_mode(self): + src = object() + dst = object() + mode_attr = object() + + with patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "_tfillpad_mode_attr", return_value=mode_attr) as build_mode, \ + patch.object(_ops._pto, "tfillpad") as tfillpad: + pto.tile.fillpad(src, dst, mode="expand") + + build_mode.assert_called_once_with("expand") + tfillpad.assert_called_once_with(src, dst, mode=mode_attr) + + def test_tile_fillpad_rejects_unknown_mode(self): + with self.assertRaisesRegex(ValueError, "normal.*in_place.*expand"): + _ops._tfillpad_mode_attr("automatic") + def test_sync_flag_names_are_exposed_without_legacy_aliases(self): preferred_names = [ "set_cross_flag", "wait_cross_flag", diff --git a/python/pto/dialects/pto.py b/python/pto/dialects/pto.py index a7b1fa44fb..034127200b 100644 --- a/python/pto/dialects/pto.py +++ b/python/pto/dialects/pto.py @@ -72,6 +72,8 @@ def _export_optional_cext_symbol(name): PadValueAttr = _pto_mod.PadValueAttr CompactMode = _pto_mod.CompactMode CompactModeAttr = _pto_mod.CompactModeAttr +TFillPadMode = _pto_mod.TFillPadMode +TFillPadModeAttr = _pto_mod.TFillPadModeAttr AccToVecMode = _pto_mod.AccToVecMode AccToVecModeAttr = _pto_mod.AccToVecModeAttr TInsertMode = _pto_mod.TInsertMode @@ -251,6 +253,8 @@ def fence_scope_attr_builder(value, context=None): "PadValueAttr", "CompactMode", "CompactModeAttr", + "TFillPadMode", + "TFillPadModeAttr", "AccToVecMode", "AccToVecModeAttr", "TInsertMode", diff --git a/test/lit/pto/fillpad_tile_native.pto b/test/lit/pto/fillpad_tile_native.pto index dc24cacc0f..434bb11a4a 100644 --- a/test/lit/pto/fillpad_tile_native.pto +++ b/test/lit/pto/fillpad_tile_native.pto @@ -22,8 +22,9 @@ module { func.func private @tfillpad_inplace_arg( %tile: !pto.tile_buf) { - pto.tfillpad_inplace ins(%tile : !pto.tile_buf) - outs(%tile : !pto.tile_buf) + pto.tfillpad ins(%tile : !pto.tile_buf) + outs(%tile : !pto.tile_buf) + {mode = #pto.tfillpad_mode} return } } @@ -31,10 +32,11 @@ module { // NATIVE-LABEL: func.func private @tfillpad_arg( // NATIVE: pto.tfillpad ins(%arg0 // NATIVE-LABEL: func.func private @tfillpad_inplace_arg( -// NATIVE: pto.tfillpad_inplace ins(%arg0 +// NATIVE: pto.tfillpad ins(%arg0 +// NATIVE-SAME: mode = #pto.tfillpad_mode // NATIVE-NOT: memref< // EMITC-LABEL: tfillpad_arg( // EMITC: TFILLPAD( // EMITC-LABEL: tfillpad_inplace_arg( -// EMITC: TFILLPAD_INPLACE( +// EMITC: TFILLPAD( diff --git a/test/lit/pto/movement_metadata_tile_native.pto b/test/lit/pto/movement_metadata_tile_native.pto index 7120bb3178..f8ee39c914 100644 --- a/test/lit/pto/movement_metadata_tile_native.pto +++ b/test/lit/pto/movement_metadata_tile_native.pto @@ -11,7 +11,7 @@ module { func.func private @tfillpad_expand_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { - pto.tfillpad_expand ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) + pto.tfillpad ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) {mode = #pto.tfillpad_mode} return } func.func private @tget_scale_addr_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { @@ -21,12 +21,13 @@ module { } // NATIVE-LABEL: @tfillpad_expand_arg -// NATIVE: pto.tfillpad_expand +// NATIVE: pto.tfillpad +// NATIVE-SAME: mode = #pto.tfillpad_mode // NATIVE-LABEL: @tget_scale_addr_arg // NATIVE: pto.tget_scale_addr // NATIVE-NOT: memref< // EMITC-LABEL: tfillpad_expand_arg( -// EMITC: TFILLPAD_EXPAND( +// EMITC: TFILLPAD( // EMITC-LABEL: tget_scale_addr_arg( // EMITC: GetScaleAddr diff --git a/test/lit/pto/tfillpad_inplace_alias_lowering.pto b/test/lit/pto/tfillpad_inplace_alias_lowering.pto index 24a583177e..e5a059766c 100644 --- a/test/lit/pto/tfillpad_inplace_alias_lowering.pto +++ b/test/lit/pto/tfillpad_inplace_alias_lowering.pto @@ -3,12 +3,12 @@ module { func.func @tfillpad_inplace_alias() { %tile = pto.alloc_tile : !pto.tile_buf - pto.tfillpad_inplace ins(%tile : !pto.tile_buf) - outs(%tile : !pto.tile_buf) + pto.tfillpad ins(%tile : !pto.tile_buf) + outs(%tile : !pto.tile_buf) + {mode = #pto.tfillpad_mode} return } } // CHECK-LABEL: AICORE void tfillpad_inplace_alias( -// CHECK: TFILLPAD_INPLACE( -// CHECK-NOT: TFILLPAD_EXPAND( +// CHECK: TFILLPAD( diff --git a/test/lit/pto/tfillpad_non_normal_mat_invalid.pto b/test/lit/pto/tfillpad_non_normal_mat_invalid.pto new file mode 100644 index 0000000000..8b4b52e454 --- /dev/null +++ b/test/lit/pto/tfillpad_non_normal_mat_invalid.pto @@ -0,0 +1,14 @@ +// RUN: not ptoas --pto-arch=a3 %s -o /dev/null 2>&1 | FileCheck %s + +module { + func.func @tfillpad_expand_mat_invalid( + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tfillpad ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {mode = #pto.tfillpad_mode} + return + } +} + +// CHECK: error: expects non-normal TFILLPAD mode only for loc=vec diff --git a/test/lit/pto/tfillpad_same_ssa_lowers_to_tfillpad.pto b/test/lit/pto/tfillpad_same_ssa_lowers_to_tfillpad.pto index 829c444030..80dd84f991 100644 --- a/test/lit/pto/tfillpad_same_ssa_lowers_to_tfillpad.pto +++ b/test/lit/pto/tfillpad_same_ssa_lowers_to_tfillpad.pto @@ -11,4 +11,4 @@ module { // CHECK-LABEL: AICORE void tfillpad_same_ssa( // CHECK: TFILLPAD( -// CHECK-NOT: TFILLPAD_INPLACE( +// CHECK-NOT: TFillPadMode::InPlace diff --git a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_expand.pto b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_expand.pto index d17af00fbf..4e0d315640 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_expand.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_expand.pto @@ -7,16 +7,16 @@ // See LICENSE in the root of the software repository for the full text of the License. // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline -// expands pto.tfillpad_expand via the PTODSL TileLib template +// expands pto.tfillpad in expand mode via the PTODSL TileLib template // // Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s // After the full tile-op-expand path on the VPTO backend, the original -// pto.tfillpad_expand should be lowered to vector-style VPTO IR. +// pto.tfillpad in expand mode should be lowered to vector-style VPTO IR. // CHECK: func.func @TFILLPAD_EXPAND -// CHECK-NOT: pto.tfillpad_expand ins +// CHECK-NOT: pto.tfillpad ins // CHECK: pto.vecscope // CHECK: pto.castptr // CHECK-DAG: pto.vdup @@ -34,10 +34,11 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { : !pto.tile_buf - pto.tfillpad_expand ins(%src : !pto.tile_buf) - outs(%dst : !pto.tile_buf) + pto.tfillpad ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {mode = #pto.tfillpad_mode} return } } diff --git a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto index a82feefce6..93ddc14923 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto @@ -31,11 +31,12 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { : !pto.tile_buf - // src 和 dst 相同，表示原地填充 padding + // src 和 dst 相同，并显式选择原地填充模式 pto.tfillpad ins(%tile : !pto.tile_buf) outs(%tile : !pto.tile_buf) + {mode = #pto.tfillpad_mode} return } } diff --git a/test/npu_validation/scripts/run_remote_npu_validation.sh b/test/npu_validation/scripts/run_remote_npu_validation.sh index 89e0995b84..450f8f4b95 100644 --- a/test/npu_validation/scripts/run_remote_npu_validation.sh +++ b/test/npu_validation/scripts/run_remote_npu_validation.sh @@ -14,7 +14,7 @@ RUN_MODE="${RUN_MODE:-npu}" # npu|sim SOC_VERSION="${SOC_VERSION:-Ascend910}" GOLDEN_MODE="${GOLDEN_MODE:-npu}" # sim|npu|skip PTO_ISA_REPO="${PTO_ISA_REPO:-https://gitcode.com/cann/pto-isa.git}" -PTO_ISA_COMMIT="${PTO_ISA_COMMIT:-ce3262e3825a235f951917eeada30e52910b6a84}" +PTO_ISA_COMMIT="${PTO_ISA_COMMIT:-27386d906e8fdcbd93aec84197939bc0b2c6caea}" DEVICE_ID="${DEVICE_ID:-0}" SKIP_CASES="${SKIP_CASES:-}" # comma/space separated testcase names RUN_ONLY_CASES="${RUN_ONLY_CASES:-}" # comma/space separated testcase names diff --git a/test/samples/Fillpad/fillpad_expand.py b/test/samples/Fillpad/fillpad_expand.py index 1bfd10011e..e064122ae8 100644 --- a/test/samples/Fillpad/fillpad_expand.py +++ b/test/samples/Fillpad/fillpad_expand.py @@ -28,6 +28,7 @@ def build(): bl = pto.BLayoutAttr.get(pto.BLayout.RowMajor, ctx) sl = pto.SLayoutAttr.get(pto.SLayout.NoneBox, ctx) pd = pto.PadValueAttr.get(pto.PadValue.Zero, ctx) + mode = pto.TFillPadModeAttr.get(pto.TFillPadMode.Expand, ctx) fractal_ab_size = pto.TileConfig.fractalABSize cfg = pto.TileBufConfigAttr.get(bl, sl, fractal_ab_size, pd, ctx) @@ -62,7 +63,7 @@ def build(): dst_tb = pto.AllocTileOp(tile_buf_32_32).result pto.TLoadOp(None, src_sv, src_tb) - pto.TFillPadExpandOp(src_tb, dst_tb) + pto.TFillPadOp(src_tb, dst_tb, mode=mode) pto.TStoreOp(None, dst_tb, dst_sv) func.ReturnOp([]) diff --git a/test/samples/Fillpad/fillpad_expand_invalid.py b/test/samples/Fillpad/fillpad_expand_invalid.py index 98fdf7206b..2f17f07375 100644 --- a/test/samples/Fillpad/fillpad_expand_invalid.py +++ b/test/samples/Fillpad/fillpad_expand_invalid.py @@ -23,6 +23,7 @@ def build(): bl = pto.BLayoutAttr.get(pto.BLayout.RowMajor, ctx) sl = pto.SLayoutAttr.get(pto.SLayout.NoneBox, ctx) pd = pto.PadValueAttr.get(pto.PadValue.Zero, ctx) + mode = pto.TFillPadModeAttr.get(pto.TFillPadMode.Expand, ctx) fractal_ab_size = pto.TileConfig.fractalABSize cfg = pto.TileBufConfigAttr.get(bl, sl, fractal_ab_size, pd, ctx) @@ -38,7 +39,7 @@ def build(): with InsertionPoint(entry): src = pto.AllocTileOp(src_ty).result dst = pto.AllocTileOp(dst_ty).result - pto.TFillPadExpandOp(src, dst) + pto.TFillPadOp(src, dst, mode=mode) func.ReturnOp([]) ok = m.operation.verify() diff --git a/test/samples/Fillpad/fillpad_expand_pad_null_invalid.py b/test/samples/Fillpad/fillpad_expand_pad_null_invalid.py index a3f250b95f..d194a093dc 100644 --- a/test/samples/Fillpad/fillpad_expand_pad_null_invalid.py +++ b/test/samples/Fillpad/fillpad_expand_pad_null_invalid.py @@ -24,6 +24,7 @@ def build(): sl = pto.SLayoutAttr.get(pto.SLayout.NoneBox, ctx) src_pd = pto.PadValueAttr.get(pto.PadValue.Zero, ctx) dst_pd = pto.PadValueAttr.get(pto.PadValue.Null, ctx) + mode = pto.TFillPadModeAttr.get(pto.TFillPadMode.Expand, ctx) fractal_ab_size = pto.TileConfig.fractalABSize src_cfg = pto.TileBufConfigAttr.get(bl, sl, fractal_ab_size, src_pd, ctx) @@ -40,7 +41,7 @@ def build(): with InsertionPoint(entry): src = pto.AllocTileOp(src_ty).result dst = pto.AllocTileOp(dst_ty).result - pto.TFillPadExpandOp(src, dst) + pto.TFillPadOp(src, dst, mode=mode) func.ReturnOp([]) ok = m.operation.verify() diff --git a/test/samples/Fillpad/fillpad_inplace.py b/test/samples/Fillpad/fillpad_inplace.py index b4159ef842..e5f0dffd5c 100644 --- a/test/samples/Fillpad/fillpad_inplace.py +++ b/test/samples/Fillpad/fillpad_inplace.py @@ -27,6 +27,7 @@ def build(): bl = pto.BLayoutAttr.get(pto.BLayout.RowMajor, ctx) sl = pto.SLayoutAttr.get(pto.SLayout.NoneBox, ctx) pd = pto.PadValueAttr.get(pto.PadValue.Zero, ctx) + mode = pto.TFillPadModeAttr.get(pto.TFillPadMode.InPlace, ctx) cfg = pto.TileBufConfigAttr.get(bl, sl, pto.TileConfig.fractalABSize, pd, ctx) tile_ty = pto.TileBufType.get([32, 32], f32, vec, [32, 32], cfg, ctx) @@ -50,7 +51,7 @@ def build(): tile = pto.AllocTileOp(tile_ty).result pto.TLoadOp(None, sv0, tile) - pto.TFillPadInplaceOp(tile, tile) + pto.TFillPadOp(tile, tile, mode=mode) pto.TStoreOp(None, tile, sv1) func.ReturnOp([]) diff --git a/test/samples/runop.sh b/test/samples/runop.sh index 967080490c..186ba63f2f 100755 --- a/test/samples/runop.sh +++ b/test/samples/runop.sh @@ -993,21 +993,16 @@ PY overall=1 continue fi - if grep -Fq "TFILLPAD_EXPAND(" "$cpp"; then - echo -e "${A}(${base}.py)\tFAIL\tpto.tfillpad should not lower via TFILLPAD_EXPAND()" + if grep -Fq "TFillPadMode::" "$cpp"; then + echo -e "${A}(${base}.py)\tFAIL\tnormal pto.tfillpad should use the default ISA mode" overall=1 continue fi fi if [[ "$base" == "fillpad_expand" ]]; then - if ! grep -Fq "TFILLPAD_EXPAND(" "$cpp"; then - echo -e "${A}(${base}.py)\tFAIL\tmissing TFILLPAD_EXPAND() lowering for pto.tfillpad_expand" - overall=1 - continue - fi - if grep -Fq "TFILLPAD(" "$cpp"; then - echo -e "${A}(${base}.py)\tFAIL\tpto.tfillpad_expand should not lower via TFILLPAD()" + if ! grep -Fq "TFILLPAD" "$cpp"; then + echo -e "${A}(${base}.py)\tFAIL\tmissing TFILLPAD<...Expand> lowering" overall=1 continue fi @@ -1032,13 +1027,8 @@ PY fi if [[ "$base" == "fillpad_inplace" ]]; then - if ! grep -Fq "TFILLPAD_INPLACE(" "$cpp"; then - echo -e "${A}(${base}.py)\tFAIL\tmissing TFILLPAD_INPLACE() lowering for pto.tfillpad_inplace" - overall=1 - continue - fi - if grep -Fq "TFILLPAD_EXPAND(" "$cpp"; then - echo -e "${A}(${base}.py)\tFAIL\tpto.tfillpad_inplace should not lower via TFILLPAD_EXPAND()" + if ! grep -Fq "TFILLPAD" "$cpp"; then + echo -e "${A}(${base}.py)\tFAIL\tmissing TFILLPAD<...InPlace> lowering" overall=1 continue fi diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_expand/tfillpad_expand.pto b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_expand/tfillpad_expand.pto index cb03e4aeda..3de41b5a0a 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_expand/tfillpad_expand.pto +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_expand/tfillpad_expand.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// TileLang ST kernels for pto.tfillpad_expand: copy src to dst and fill padding. +// TileLang ST kernels for pto.tfillpad in expand mode: copy src to dst and fill padding. // Matches C++ test cases: case 8, 9 // Compiled by ptoas --enable-insert-sync --enable-tile-op-expand --pto-backend=vpto // @@ -56,8 +56,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) outs(%src : !pto.tile_buf) - pto.tfillpad_expand ins(%src : !pto.tile_buf) - outs(%dst : !pto.tile_buf) + pto.tfillpad ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {mode = #pto.tfillpad_mode} pto.tstore ins(%dst : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x64x16xui16>) @@ -108,8 +109,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) outs(%src : !pto.tile_buf) - pto.tfillpad_expand ins(%src : !pto.tile_buf) - outs(%dst : !pto.tile_buf) + pto.tfillpad ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {mode = #pto.tfillpad_mode} pto.tstore ins(%dst : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x32xui16>) diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto index 0d80ee5595..4421609d6d 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto @@ -54,8 +54,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) outs(%tile_buf : !pto.tile_buf) - pto.tfillpad_inplace ins(%tile_buf : !pto.tile_buf) - outs(%tile_buf : !pto.tile_buf) + pto.tfillpad ins(%tile_buf : !pto.tile_buf) + outs(%tile_buf : !pto.tile_buf) + {mode = #pto.tfillpad_mode} pto.tstore ins(%tile_buf : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x64x16xf32>) @@ -103,9 +104,10 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) outs(%tile_buf : !pto.tile_buf) - // tfillpad_inplace: src_valid == dst_valid, no expansion - pto.tfillpad_inplace ins(%tile_buf : !pto.tile_buf) - outs(%tile_buf : !pto.tile_buf) + // tfillpad in_place: src_valid == dst_valid, no expansion + pto.tfillpad ins(%tile_buf : !pto.tile_buf) + outs(%tile_buf : !pto.tile_buf) + {mode = #pto.tfillpad_mode} // Store full tile pto.tstore ins(%tile_buf : !pto.tile_buf) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_expand/tfillpad_expand.pto b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_expand/tfillpad_expand.pto index 6e68ba00be..3d68127035 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_expand/tfillpad_expand.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_expand/tfillpad_expand.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// TileLang ST kernels for pto.tfillpad_expand: copy src to dst and fill padding. +// TileLang ST kernels for pto.tfillpad in expand mode: copy src to dst and fill padding. // Matches C++ test cases: case 8, 9 // Compiled by ptoas --enable-insert-sync --enable-tile-op-expand --pto-backend=vpto // @@ -60,8 +60,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) outs(%src : !pto.tile_buf) - pto.tfillpad_expand ins(%src : !pto.tile_buf) - outs(%dst : !pto.tile_buf) + pto.tfillpad ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {mode = #pto.tfillpad_mode} pto.tstore ins(%dst : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x32xui16>) @@ -111,11 +112,12 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) outs(%src : !pto.tile_buf) - pto.tfillpad_expand ins(%src : !pto.tile_buf) - outs(%dst : !pto.tile_buf) + pto.tfillpad ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {mode = #pto.tfillpad_mode} pto.tstore ins(%dst : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x64xi8>) return } -} \ No newline at end of file +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto index f4d606df35..36cf1dff95 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto @@ -59,13 +59,14 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) outs(%tile_buf : !pto.tile_buf) - // tfillpad_inplace: src_valid == dst_valid, no expansion - pto.tfillpad_inplace ins(%tile_buf : !pto.tile_buf) - outs(%tile_buf : !pto.tile_buf) + // tfillpad in_place: src_valid == dst_valid, no expansion + pto.tfillpad ins(%tile_buf : !pto.tile_buf) + outs(%tile_buf : !pto.tile_buf) + {mode = #pto.tfillpad_mode} // Store full tile pto.tstore ins(%tile_buf : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x16xf32>) return } -} \ No newline at end of file +} diff --git a/tools/ptobc/generated/ptobc_opcodes_v0.h b/tools/ptobc/generated/ptobc_opcodes_v0.h index cce943b3bc..ec0a1df46f 100644 --- a/tools/ptobc/generated/ptobc_opcodes_v0.h +++ b/tools/ptobc/generated/ptobc_opcodes_v0.h @@ -97,8 +97,6 @@ inline constexpr OpInfo kOpTable[] = { {0x1021, "pto.textract", 0, 0x00, 0x00, 4, 0, 0, 0x00}, {0x1022, "pto.textract_fp", 0, 0x00, 0x00, 5, 0, 0, 0x00}, {0x1023, "pto.tfillpad", 0, 0x00, 0x00, 2, 0, 0, 0x00}, - {0x1024, "pto.tfillpad_expand", 0, 0x00, 0x00, 2, 0, 0, 0x00}, - {0x1025, "pto.tfillpad_inplace", 0, 0x00, 0x00, 2, 0, 0, 0x00}, {0x1026, "pto.tfmod", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x1027, "pto.tfmods", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x1028, "pto.tgather", 0, 0x00, 0x02, 0, 0, 0, 0x00}, @@ -306,8 +304,6 @@ inline std::optional lookupOpcodeByName(llvm::StringRef name) { .Case("pto.textract", 0x1021) .Case("pto.textract_fp", 0x1022) .Case("pto.tfillpad", 0x1023) - .Case("pto.tfillpad_expand", 0x1024) - .Case("pto.tfillpad_inplace", 0x1025) .Case("pto.tfmod", 0x1026) .Case("pto.tfmods", 0x1027) .Case("pto.tgather", 0x1028) @@ -502,8 +498,6 @@ inline std::optional lookupOpcodeAndVariantByFullName(llvm::St .Case("pto.textract", OpcodeAndVariant{0x1021, 0, 0}) .Case("pto.textract_fp", OpcodeAndVariant{0x1022, 0, 0}) .Case("pto.tfillpad", OpcodeAndVariant{0x1023, 0, 0}) - .Case("pto.tfillpad_expand", OpcodeAndVariant{0x1024, 0, 0}) - .Case("pto.tfillpad_inplace", OpcodeAndVariant{0x1025, 0, 0}) .Case("pto.tfmod", OpcodeAndVariant{0x1026, 0, 0}) .Case("pto.tfmods", OpcodeAndVariant{0x1027, 0, 0}) .Case("pto.tgather", OpcodeAndVariant{0x1028, 0, 0}) From 3d8482a20a70c1b72417ab0c335308023b61e012 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Tue, 4 Aug 2026 13:16:56 +0800 Subject: [PATCH 003/122] feat: unify fp operand forms --- docs/PTO_IR_manual.md | 118 +--- ...ptoas-tile-native-mainline-op-migration.md | 9 +- include/PTO/IR/PTOOps.td | 147 +---- lib/PTO/IR/PTO.cpp | 504 +++--------------- lib/PTO/Transforms/ConvertToPTOOp.cpp | 4 +- lib/PTO/Transforms/PTOToEmitC.cpp | 217 ++------ lib/TileOps/__init__.py | 4 +- lib/TileOps/a5/textract_fp.py | 32 +- lib/TileOps/a5/tstore.py | 6 +- ptodsl/ptodsl/_ops.py | 33 +- ptodsl/ptodsl/_tile_namespace.py | 6 +- ptodsl/tests/test_tilelib_catalog.py | 62 +-- ptodsl/tests/test_vector_cube_ops.py | 23 + test/lit/pto/cube_tile_ops_positive.pto | 2 +- .../lit/pto/extract_insert_fp_tile_native.pto | 10 +- test/lit/pto/store_fp_tile_native.pto | 4 +- test/lit/pto/textract_acc_to_vec_a5_emitc.pto | 2 +- test/lit/pto/textract_forms_emitc.pto | 2 +- test/lit/pto/textract_fp_a3_lowering.pto | 2 +- ...extract_tinsert_low_precision_a5_emitc.pto | 4 +- test/lit/pto/tinsert_a5_modes_emitc.pto | 2 +- test/lit/pto/tinsert_forms_emitc.pto | 2 +- test/lit/pto/tinsert_fp_a3_lowering.pto | 2 +- .../tinsert_fp_scaling_tile_role_emitc.pto | 2 +- test/lit/pto/tmov_fp_tile_native.pto | 5 +- .../lit/pto/tstore_fp_insert_sync_effects.pto | 10 +- test/lit/pto/tstore_fp_invalid_dtype.pto | 4 +- ...d_tile_op_tilelang_textract_fp_acc2mat.pto | 10 +- ...ile_op_tilelang_textract_fp_all_dtypes.pto | 32 +- .../textract_fp_verify_invalid_dst_loc.pto | 10 +- .../textract_fp_verify_invalid_dtype_pair.pto | 10 +- .../textract_fp_verify_invalid_fp_loc.pto | 8 +- .../textract_fp_verify_invalid_src_layout.pto | 10 +- .../textract_fp_verify_invalid_src_loc.pto | 10 +- test/samples/Extract/extract_fp.py | 2 +- test/samples/Movfp/movfp.py | 2 +- test/samples/Storefp/storefp.py | 2 +- test/samples/Storefp/storefp_invalid.py | 2 +- test/samples/TInsert/tinsert_fp.py | 2 +- test/samples/runop.sh | 4 +- .../testcase/textract_fp/textract_fp.pto | 4 +- .../st/testcase/textract_fp/textract_fp.pto | 4 +- .../testcase/tstore_acc2gm/tstore_acc2gm.pto | 4 +- tools/ptobc/generated/ptobc_opcodes_v0.h | 18 +- .../ptobc/testdata/tstore_fp_v0_roundtrip.pto | 2 +- tools/ptobc/tests/tstore_fp_v0_encode.sh | 3 +- 46 files changed, 344 insertions(+), 1013 deletions(-) diff --git a/docs/PTO_IR_manual.md b/docs/PTO_IR_manual.md index 2fb7e36efa..d884d9d8e8 100644 --- a/docs/PTO_IR_manual.md +++ b/docs/PTO_IR_manual.md @@ -1320,7 +1320,7 @@ Lowering maps `%ctx` to `pto::PrefetchAsyncContext`, emits ##### `pto.tstore` - Store Tile to Partition View -**Summary:** Stores a 2-D tile buffer back to a 2-D partition view. Supports phase/atomic/relu/pre-quant controls that lower to the corresponding `TSTORE` template overload family. +**Summary:** Stores a 2-D tile buffer back to a 2-D partition view. Supports phase/atomic/relu/pre-quant controls and an optional scaling tile. The scaling-tile form lowers to `TSTORE_FP`. **Semantics:** @@ -1335,6 +1335,7 @@ For each element (i, j) in the tile valid region: |------|------|---------|-------------| | `src` | `pto.tile_buf` | `NA` |Source tile buffer | | `dst` | `PartitionTensorViewType` | `NA` | Destination partition view | +| `fp` | `pto.tile_buf` (optional) | `NA` | Scaling tile (`loc=scaling`) for accumulator conversion | | `preQuantScalar` | `i64` (optional) | `NA` |Optional scalar used by pre-quantized `acc` store forms | | `stPhase` | `#pto` | `unspecified` | Store phase selector (`unspecified/partial/final`) | | `atomicType` | `#pto` | `atomic_none` | Atomic mode (`atomic_none/atomic_add`) | @@ -1348,7 +1349,9 @@ For each element (i, j) in the tile valid region: - `src` must be `!pto.tile_buf`, `dst` must be `!pto.partition_tensor_view`. - Static `dst` shape dims must be positive, and static `src` valid-shape dims must be non-negative. - - If `preQuantScalar` is present, `src` must be `loc=acc`. + - `fp` and `preQuantScalar` are mutually exclusive. + - If `fp` or `preQuantScalar` is present, `src` must be `loc=acc`. + - `fp` must use `loc=scaling`; the fp form uses the default `stPhase`. - If `reluPreMode != no_relu`, `src` must be `loc=acc`. - A2/A3 checks: - `src.loc` must be one of `vec/mat/acc`. @@ -1615,6 +1618,9 @@ For each element (i, j): | `src` | `pto.tile_buf` | Source tile | | `tmp` | `pto.tile_buf` | Temporary workspace operand required by the current DPS form | | `dst` | `pto.tile_buf` | Destination tile | +| `fp` | `pto.tile_buf` | Optional scaling tile (`loc=scaling`) for accumulator conversion | +| `preQuantScalar` | `i64` | Optional scalar pre-quant parameter | +| `accToVecMode` | `pto.acc_to_vec_mode` | Optional A5 acc-to-vec mode | **Results:** None. Writes into `dst` via DPS pattern. @@ -7976,7 +7982,9 @@ dst[i + indexRow, j + indexCol] = src[i, j] **Hardware Mapping:** -- Lowers to **`TINSERT(dst, src, indexRow, indexCol)`** +- Without `fp`, lowers to **`TINSERT(dst, src, indexRow, indexCol)`**. +- With `fp`, lowers to **`TINSERT_FP(dst, src, fp, indexRow, indexCol)`**; + an explicit A5 `accToVecMode` selects the fp-parameterized `TINSERT` overload. - Uses the target data-movement pipeline: `Vec -> Vec` uses `PIPE_V`, A5 `Vec -> Mat` uses `PIPE_MTE3`, and regular `Acc -> Mat` uses `PIPE_FIX`. @@ -7984,6 +7992,7 @@ dst[i + indexRow, j + indexCol] = src[i, j] ```mlir pto.tinsert ins(%src, %row, %col : !pto.tile_buf<...>, index, index) outs(%dst : !pto.tile_buf<...>) +pto.tinsert ins(%src, %row, %col : !pto.tile_buf<...>, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf<...>) ``` --- @@ -8006,6 +8015,9 @@ dst[i, j] = src[i + indexRow, j + indexCol] | `indexRow` | `Index` | Starting row | | `indexCol` | `Index` | Starting column | | `dst` | `pto.tile_buf` | Destination tile | +| `fp` | `pto.tile_buf` | Optional scaling tile (`loc=scaling`) for accumulator conversion | +| `preQuantScalar` | `i64` | Optional scalar pre-quant parameter | +| `accToVecMode` | `pto.acc_to_vec_mode` | Optional A5 acc-to-vec mode | **Results:** None. Writes into `dst` via DPS pattern. @@ -8026,12 +8038,15 @@ dst[i, j] = src[i + indexRow, j + indexCol] **Hardware Mapping:** -- Executes on the **Vector pipeline** (`PIPE_V`) +- Base forms lower to `TEXTRACT`; an `fp` form without mode lowers to + `TEXTRACT_FP`, while an explicit A5 `accToVecMode` selects the + fp-parameterized `TEXTRACT` overload. **Basic Example:** ```mlir -pto.textract ins(%src[%row, %col] : !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) +pto.textract ins(%src, %row, %col : !pto.tile_buf<...>, index, index) outs(%dst : !pto.tile_buf<...>) +pto.textract ins(%src, %row, %col : !pto.tile_buf<...>, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf<...>) ``` --- @@ -8484,54 +8499,6 @@ pto.tget_scale_addr ins(%src : !pto.tile_buf, !pto.tile_buf<...>) - outs(%dst : !pto.tile_buf<...>) -``` - ---- - ##### `pto.tquant` - Quantize Tile with Scaling Tile **Summary:** Quantizes `f32` source tile elements into a lower-precision integer format using a scaling (`fp`) tile. The quantization mode is controlled by the `quant_type` attribute. @@ -8647,51 +8614,6 @@ pto.tquant.mx ins(%src : !pto.tile_buf<...>) --- -##### `pto.tstore_fp` - Store Accumulator with Scaling - -**Summary:** Stores an accumulator tile into global memory using a scaling (`fp`) tile. - -**Semantics:** - -``` -dst[...] = Convert(src[i, j]; fp) -``` - -**Arguments:** - -| Name | Type | Description | -|------|------|-------------| -| `src` | `pto.tile_buf` | Source accumulator tile | -| `fp` | `pto.tile_buf` | Scaling tile | -| `dst` | `PartitionTensorViewType` | Destination memory | - -**Results:** None. Writes into `dst` via DPS pattern. - -**Constraints & Verification:** - -- **Implementation checks (A2A3)** - - Source TileType only suport `loc==acc` - - Source dtype must be `i32` or `f32`. - - Shape constraints: `1 <= cols <= 4095`; - - Runtime: `1 <= src valid column <= 4095`. - - `fp` is used to configure scaling/FPC state; no separate PTO-visible static constraint is enforced on its shape. -- **Implementation checks (A5)** - - Source TileType only suport `loc==acc` - - `fp` is used to configure scaling/FPC state; no separate PTO-visible static constraint is enforced on its shape. - -**Hardware Mapping:** - -- Executes on the **DMA pipeline** (`PIPE_MTE3`) - -**Basic Example:** - -```mlir -pto.tstore_fp ins(%acc, %fp : !pto.tile_buf<...>, !pto.tile_buf<...>) - outs(%dst : memref<...>) -``` - ---- - ### 4.16 Synchronization Operations ##### `pto.barrier` diff --git a/docs/designs/ptoas-tile-native-mainline-op-migration.md b/docs/designs/ptoas-tile-native-mainline-op-migration.md index f76d019dff..bfa5e3050c 100644 --- a/docs/designs/ptoas-tile-native-mainline-op-migration.md +++ b/docs/designs/ptoas-tile-native-mainline-op-migration.md @@ -161,8 +161,7 @@ PTO tile/view IR | `pto.tgemv_mx` | MX scale role | | `pto.tgemv_mx_acc` | MX scale + accumulator effects | | `pto.tgemv_mx_bias` | MX scale + bias effects | -| `pto.tmov` | identity removal、view metadata 和 src/dst alias | -| `pto.tmov_fp` | fp/pre-quant optional operands和地址空间 | +| `pto.tmov` | identity removal、view metadata、src/dst alias，以及 fp/pre-quant optional operands 和地址空间 | | `pto.tabs` | unary Read(src)/Write(dst) | | `pto.tand` | binary input/output effects | | `pto.tands` | scalar form operand 顺序 | @@ -186,10 +185,8 @@ PTO tile/view IR | `pto.tdiv` | precision attr 和 inplace policy | | `pto.tdivs` | scalar 顺序和 precision attr | | `pto.texpands` | shape扩展和 scalar operand | -| `pto.textract` | tile role、offset 和可选 pre-quant | -| `pto.textract_fp` | fp tile role和地址空间 | -| `pto.tinsert` | materialize pass 当前对 tile config 有特殊推断；目标是由 result type 完整携带 | -| `pto.tinsert_fp` | fp/pre-quant tile role | +| `pto.textract` | tile role、offset、fp tile 地址空间和可选 pre-quant | +| `pto.tinsert` | materialize pass 当前对 tile config 有特殊推断；包含 fp/pre-quant tile role，目标是由 result type 完整携带 | | `pto.tfillpad` | 显式 mode、src/dst alias、MemoryEffects 和 A5 MAT/PIPE 选择 | | `pto.tsetval` | tile writer 和 result type | | `pto.tgetval` | tile reader和 scalar result | diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 5c29a5f9a6..aeaf3a33d3 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -786,6 +786,7 @@ def TPrefetchAsyncOp : PTO_Op<"tprefetch_async"> { } def TStoreOp: PTO_TOp<"tstore", [ + AttrSizedOperandSegments, PTO_DpsInitOpInterface, OpPipeInterface, DeclareOpInterfaceMethods @@ -801,6 +802,7 @@ def TStoreOp: PTO_TOp<"tstore", [ let arguments = (ins PTODpsType:$src, PTODpsType:$dst, + Optional:$fp, Optional:$preQuantScalar, DefaultValuedAttr:$stPhase, DefaultValuedAttr:$atomicType, @@ -815,6 +817,7 @@ def TStoreOp: PTO_TOp<"tstore", [ let assemblyFormat = [{ `ins` `(` $src `:` qualified(type($src)) + (`fp` $fp^ `:` qualified(type($fp)))? (`,` $preQuantScalar^ `:` type($preQuantScalar))? `)` `outs` `(` $dst `:` qualified(type($dst) ) `)` attr-dict @@ -4512,6 +4515,7 @@ def TExpandsOp : PTO_TOp<"texpands", [ } def TExtractOp : PTO_TOp<"textract", [ + AttrSizedOperandSegments, PTO_DpsInitOpInterface, OpPipeInterface, DeclareOpInterfaceMethods @@ -4522,8 +4526,9 @@ def TExtractOp : PTO_TOp<"textract", [ PTODpsType:$src, Index:$indexRow, Index:$indexCol, - Optional:$preQuantScalar, PTODpsType:$dst, + Optional:$fp, + Optional:$preQuantScalar, OptionalAttr:$accToVecMode, DefaultValuedAttr:$reluPreMode ); @@ -4533,7 +4538,10 @@ def TExtractOp : PTO_TOp<"textract", [ let hasVerifier = 1; let assemblyFormat = [{ - `ins` `(` $src `,` $indexRow `,` $indexCol (`,` $preQuantScalar^)? `:` qualified(type($src)) `,` type($indexRow) `,` type($indexCol) (`,` type($preQuantScalar)^)? `)` + `ins` `(` $src `,` $indexRow `,` $indexCol + (`,` $preQuantScalar^)? `:` qualified(type($src)) `,` type($indexRow) `,` type($indexCol) + (`,` type($preQuantScalar)^)? + (`fp` $fp^ `:` qualified(type($fp)))? `)` `outs` `(` $dst `:` qualified(type($dst) ) `)` attr-dict }]; @@ -4575,7 +4583,8 @@ def TExtractOp : PTO_TOp<"textract", [ if ((s == ::mlir::pto::AddressSpace::VEC && d == ::mlir::pto::AddressSpace::MAT) || (s == ::mlir::pto::AddressSpace::ACC && - d == ::mlir::pto::AddressSpace::MAT)) { + (d == ::mlir::pto::AddressSpace::MAT || + d == ::mlir::pto::AddressSpace::VEC))) { return ::mlir::pto::PIPE::PIPE_FIX; } @@ -4591,39 +4600,6 @@ def TExtractOp : PTO_TOp<"textract", [ }]; } -def TExtractFPOp : PTO_TOp<"textract_fp", [ - PTO_DpsInitOpInterface, - OpPipeInterface, - DeclareOpInterfaceMethods -]> { - let summary = "Extract acc tile window into dst using fp/scaling tile (tilebuf, DPS)"; - - let arguments = (ins - PTODpsType:$src, - PTODpsType:$fp, - Index:$indexRow, - Index:$indexCol, - PTODpsType:$dst, - OptionalAttr:$accToVecMode, - DefaultValuedAttr:$reluPreMode - ); - - let results = (outs); - - let hasVerifier = 1; - - let assemblyFormat = [{ - `ins` `(` $src `,` $fp `,` $indexRow `,` $indexCol `:` qualified(type($src)) `,` qualified(type($fp)) `,` type($indexRow) `,` type($indexCol) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; - - let extraClassDeclaration = [{ - ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_FIX; } - ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } - }]; -} - def TInsertOp : PTO_TOp<"tinsert", [ AttrSizedOperandSegments, PTO_DpsInitOpInterface, @@ -4712,39 +4688,6 @@ def TInsertOp : PTO_TOp<"tinsert", [ }]; } -def TInsertFPOp : PTO_TOp<"tinsert_fp", [ - PTO_DpsInitOpInterface, - OpPipeInterface, - DeclareOpInterfaceMethods -]> { - let summary = "Insert acc tile window into dst using fp/scaling tile (tilebuf, DPS)"; - - let arguments = (ins - PTODpsType:$src, - PTODpsType:$fp, - Index:$indexRow, - Index:$indexCol, - PTODpsType:$dst, - OptionalAttr:$accToVecMode, - DefaultValuedAttr:$reluPreMode - ); - - let results = (outs); - - let hasVerifier = 1; - - let assemblyFormat = [{ - `ins` `(` $src `,` $fp `,` $indexRow `,` $indexCol `:` qualified(type($src)) `,` qualified(type($fp)) `,` type($indexRow) `,` type($indexCol) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; - - let extraClassDeclaration = [{ - ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_FIX; } - ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } - }]; -} - def TFillPadOp : PTO_TOp<"tfillpad", [ PTO_DpsInitOpInterface, OpPipeInterface, @@ -5033,37 +4976,6 @@ def TMinSOp : PTO_TOp<"tmins", [ }]; } -def TMovFPOp : PTO_TOp<"tmov.fp", [ - PTO_DpsInitOpInterface, - OpPipeInterface, - DeclareOpInterfaceMethods -]> { - let summary = "TMOV_FP: move/convert using fp (scaling) tile (tilebuf, DPS)"; - - let arguments = (ins - PTODpsType:$src, - PTODpsType:$fp, - PTODpsType:$dst - ); - - let results = (outs); - - let hasVerifier = 1; - - let assemblyFormat = [{ - `ins` `(` $src `,` $fp `:` qualified(type($src)) `,` qualified(type($fp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; - - let extraClassDeclaration = [{ - // TMOV_FP is an ACC->MAT move (Cc->Cb) with vector quant parameters in - // SCALING (fbuf). Treat it as a data-movement op for sync insertion. - ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_MTE1; } - ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } - }]; -} - def TMrgSortOp: PTO_TOp<"tmrgsort", [ AttrSizedOperandSegments, PTO_DpsInitOpInterface, @@ -6610,41 +6522,6 @@ def TSqrtOp: PTO_TOp<"tsqrt", [ }]; } -//===----------------------------------------------------------------------===// -// PTOOps.td (add TSTORE_FP TBDPS/tile buffer op) -//===----------------------------------------------------------------------===// - -def TStoreFPOp: PTO_TOp<"tstore_fp", [ - PTO_DpsInitOpInterface, - OpPipeInterface, - DeclareOpInterfaceMethods, -]> { - let summary = "TSTORE_FP: Store an accumulator tile into global memory using a scaling (fp) tile for vector quantization parameters."; - - let arguments = (ins - PTODpsType:$src, - PTODpsType:$fp, - PTODpsType:$dst - ); - - let results = (outs); - - let hasVerifier = 1; - - let assemblyFormat = [{ - `ins` `(` $src `,` $fp `:` qualified(type($src)) `,` qualified(type($fp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; - let extraClassDeclaration = [{ - ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_FIX; } - ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } - }]; -} -//===----------------------------------------------------------------------===// -// PTOOps.td (add TSUB TBDPS/tile buffer op) -//===----------------------------------------------------------------------===// - def TSubOp: PTO_TOp<"tsub", [ PTO_DpsInitOpInterface, OpPipeInterface, diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index bf6caf9b4e..965b391d79 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -3579,6 +3579,13 @@ LogicalResult mlir::pto::SyncWaitOp::verify() { } LogicalResult TStoreOp::verify() { + const bool hasFp = static_cast(getFp()); + const bool hasPreQuant = static_cast(getPreQuantScalar()); + if (hasFp && hasPreQuant) + return emitOpError("expects fp and preQuantScalar to be mutually exclusive"); + if (hasFp && getStPhase() != pto::STPhase::Unspecified) + return emitOpError("expects fp form to use the default stPhase"); + auto verifyCommon = [&](bool allowLowPrecision) -> FailureOr> { @@ -3590,6 +3597,16 @@ LogicalResult TStoreOp::verify() { } if (failed(verifyTileBufCommon(*this, srcTile, "src", allowLowPrecision))) return failure(); + if (hasFp) { + Type fpTy = getFp().getType(); + if (failed(verifyTileBufCommon(*this, fpTy, "fp", allowLowPrecision))) + return failure(); + auto fpSpace = getPTOMemorySpaceEnum(fpTy); + if (!fpSpace || *fpSpace != pto::AddressSpace::SCALING) { + emitOpError("expects fp to use loc=scaling"); + return failure(); + } + } for (auto [idx, dim] : llvm::enumerate(dstPart.getShape())) { if (dim != ShapedType::kDynamic && dim <= 0) { emitOpError() << "expects dst shape[" << idx << "] to be positive"; @@ -3624,7 +3641,8 @@ LogicalResult TStoreOp::verify() { auto dstElemCount = getStaticElemCount(dstPart.getShape()); auto srcValidElemCount = getStaticElemCount(srcValid); - if (dstElemCount && srcValidElemCount && *dstElemCount != *srcValidElemCount) { + if (!hasFp && dstElemCount && srcValidElemCount && + *dstElemCount != *srcValidElemCount) { emitOpError() << "expects dst static element count (" << *dstElemCount << ") to match src valid_shape static element count (" << *srcValidElemCount << ")"; @@ -3638,7 +3656,6 @@ LogicalResult TStoreOp::verify() { ty.isInteger(64) || ty.isF16() || ty.isBF16() || ty.isF32(); }; auto isI8Like = [&](Type ty) -> bool { return ty.isInteger(8); }; - bool hasPreQuant = static_cast(getPreQuantScalar()); auto reluMode = getReluPreMode(); auto verifyA2A3 = [&]() -> LogicalResult { @@ -3651,16 +3668,16 @@ LogicalResult TStoreOp::verify() { *srcSpace != pto::AddressSpace::MAT && *srcSpace != pto::AddressSpace::ACC)) return emitOpError("expects A2/A3 tstore src to use loc=vec, loc=mat, or loc=acc"); - if (hasPreQuant && *srcSpace != pto::AddressSpace::ACC) - return emitOpError("expects preQuantScalar form to use loc=acc src"); + if ((hasFp || hasPreQuant) && *srcSpace != pto::AddressSpace::ACC) + return emitOpError("expects fp/preQuantScalar form to use loc=acc src"); if (reluMode != pto::ReluPreMode::NoRelu && *srcSpace != pto::AddressSpace::ACC) return emitOpError("expects reluPreMode form to use loc=acc src"); Type srcElem = srcTile.getElementType(); Type dstElem = dstPart.getElementType(); if (*srcSpace == pto::AddressSpace::VEC || *srcSpace == pto::AddressSpace::MAT) { - if (hasPreQuant) - return emitOpError("expects preQuantScalar form to use loc=acc src"); + if (hasFp || hasPreQuant) + return emitOpError("expects fp/preQuantScalar form to use loc=acc src"); if (isPTOLowPrecisionType(dstElem)) return emitOpError("expects A2/A3 vec/mat tstore low-precision dst element types to be unsupported"); if (!isLoadStoreElemType(srcElem)) @@ -3680,7 +3697,7 @@ LogicalResult TStoreOp::verify() { if (!isI8Like(dstElem)) return emitOpError("expects A2/A3 acc preQuantScalar tstore dst type to be i8/ui8"); } - } else { + } else if (!hasFp) { if (!(dstElem.isInteger(32) || dstElem.isF32() || dstElem.isF16() || dstElem.isBF16())) return emitOpError("expects A2/A3 acc tstore dst element type to be i32/f32/f16/bf16"); @@ -3706,16 +3723,16 @@ LogicalResult TStoreOp::verify() { if (!srcSpace || (*srcSpace != pto::AddressSpace::VEC && *srcSpace != pto::AddressSpace::ACC)) return emitOpError("expects A5 tstore src to use loc=vec or loc=acc"); - if (hasPreQuant && *srcSpace != pto::AddressSpace::ACC) - return emitOpError("expects preQuantScalar form to use loc=acc src"); + if ((hasFp || hasPreQuant) && *srcSpace != pto::AddressSpace::ACC) + return emitOpError("expects fp/preQuantScalar form to use loc=acc src"); if (reluMode != pto::ReluPreMode::NoRelu && *srcSpace != pto::AddressSpace::ACC) return emitOpError("expects reluPreMode form to use loc=acc src"); Type srcElem = srcTile.getElementType(); Type dstElem = dstPart.getElementType(); if (*srcSpace == pto::AddressSpace::VEC) { - if (hasPreQuant) - return emitOpError("expects preQuantScalar form to use loc=acc src"); + if (hasFp || hasPreQuant) + return emitOpError("expects fp/preQuantScalar form to use loc=acc src"); if (!isA5TLoadStoreTransferElemType(srcElem)) return emitOpError("expects A5 vec tstore src element type to be i8/i16/i32/i64/f16/bf16/f32/f8/hif8/fp4"); if (getElemByteSize(srcElem) != getElemByteSize(dstElem)) @@ -3741,7 +3758,7 @@ LogicalResult TStoreOp::verify() { if (hasPreQuant) { if (!isA5AccStorePreQuantDstType(srcElem, dstElem)) return emitOpError("expects A5 acc preQuantScalar tstore dst type to be i8/ui8/f16/bf16/f32/hif8/f8E4M3"); - } else { + } else if (!hasFp) { if (!(dstElem.isInteger(32) || dstElem.isF32() || dstElem.isF16() || dstElem.isBF16())) return emitOpError("expects A5 acc tstore dst element type to be i32/f32/f16/bf16"); @@ -6537,9 +6554,11 @@ mlir::LogicalResult mlir::pto::TExtractOp::verify() { ty.getSLayoutValueI32() == static_cast(pto::SLayout::NoneBox); }; Value preQuantScalar = getPreQuantScalar(); + Value fp = getFp(); auto reluMode = getReluPreMode(); auto accToVecModeAttr = getAccToVecModeAttr(); const bool hasPreQuantScalar = static_cast(preQuantScalar); + const bool hasFp = static_cast(fp); const bool hasRelu = reluMode != pto::ReluPreMode::NoRelu; const bool hasAccToVecMode = static_cast(accToVecModeAttr); auto verifyCommon = [&](bool allowLowPrecision) @@ -6557,11 +6576,19 @@ mlir::LogicalResult mlir::pto::TExtractOp::verify() { failed(verifyTileBufCommon(*this, dstTy, "dst", allowLowPrecision)) || failed(verifyNonNegativeIndexRowCol( *getOperation(), getIndexRow(), getIndexCol(), - /*includeIndexAndIntOpsInConstFold=*/false)) || + /*includeIndexAndIntOpsInConstFold=*/hasFp)) || failed(verifyExtractStaticBoundsCommon( *getOperation(), getIndexRow(), getIndexCol(), srcTy, dstTy, - /*includeIndexAndIntOpsInConstFold=*/false))) + /*includeIndexAndIntOpsInConstFold=*/hasFp))) return failure(); + if (hasFp) { + Type fpTy = fp.getType(); + if (failed(verifyTileBufCommon(*this, fpTy, "fp", allowLowPrecision))) + return failure(); + auto fpSpace = getPTOMemorySpaceEnum(fpTy); + if (!fpSpace || *fpSpace != pto::AddressSpace::SCALING) + return emitOpError("expects fp to use loc=scaling"); + } auto srcSpace = getPTOMemorySpaceEnum(srcTy); auto dstSpace = getPTOMemorySpaceEnum(dstTy); Type srcElem = getElemTy(srcTy); @@ -6579,8 +6606,13 @@ mlir::LogicalResult mlir::pto::TExtractOp::verify() { return failure(); auto [srcTy, dstTy, srcTb, dstTb, srcElem, dstElem, srcSpace, dstSpace] = *common; - if (!isA2A3ExtractElemType(dstElem)) + if (!isA2A3ExtractElemType(dstElem) && + !(hasFp && dstElem.isInteger(16))) return emitOpError("expects A2/A3 textract element type to be i8/f16/bf16/f32"); + if (hasFp && hasPreQuantScalar) + return emitOpError("expects fp and preQuantScalar to be mutually exclusive"); + if (hasFp && (!srcSpace || *srcSpace != pto::AddressSpace::ACC)) + return emitOpError("expects fp form to use loc=acc src"); if (hasPreQuantScalar && (!srcSpace || *srcSpace != pto::AddressSpace::ACC)) return emitOpError("expects preQuantScalar form to use loc=acc src"); if (hasRelu && (!srcSpace || *srcSpace != pto::AddressSpace::ACC)) @@ -6606,7 +6638,7 @@ mlir::LogicalResult mlir::pto::TExtractOp::verify() { return emitOpError("expects A2/A3 acc-source textract dst to use blayout=col_major and slayout=row_major"); if (dstTb.getSFractalSizeI32() != 512) return emitOpError("expects A2/A3 acc-source textract dst fractal size to be 512"); - if (hasPreQuantScalar) { + if (hasFp || hasPreQuantScalar) { if (!isA2A3AccQuantExtractTypePair(srcElem, dstElem)) return emitOpError( "expects A2/A3 acc preQuantScalar textract element types to be " @@ -6643,6 +6675,10 @@ mlir::LogicalResult mlir::pto::TExtractOp::verify() { *common; if (!isA5ExtractElemType(dstElem)) return emitOpError("expects A5 textract element type to be an fp8/f16/bf16/f32 or int8 family type"); + if (hasFp && hasPreQuantScalar) + return emitOpError("expects fp and preQuantScalar to be mutually exclusive"); + if (hasFp && (!srcSpace || *srcSpace != pto::AddressSpace::ACC)) + return emitOpError("expects fp form to use loc=acc src"); if (hasPreQuantScalar && (!srcSpace || *srcSpace != pto::AddressSpace::ACC)) return emitOpError("expects preQuantScalar form to use loc=acc src"); if (hasRelu && (!srcSpace || *srcSpace != pto::AddressSpace::ACC)) @@ -6700,7 +6736,7 @@ mlir::LogicalResult mlir::pto::TExtractOp::verify() { if (!isRowMajorNoneBoxND(dstTb)) return emitOpError("expects A5 acc->vec textract dst to use ND layout (blayout=row_major, slayout=none_box)"); } - if (hasPreQuantScalar) { + if (hasFp || hasPreQuantScalar) { if (!isA5AccQuantExtractTypePair(srcElem, dstElem)) return emitOpError( "expects A5 acc preQuantScalar textract element types to be " @@ -6803,6 +6839,9 @@ mlir::LogicalResult mlir::pto::TInsertOp::verify() { auto fpTy = getFp().getType(); auto fpTb = dyn_cast(fpTy); if (!fpTb) return emitOpError("expects fp to be !pto.tile_buf"); + if (failed(verifyTileBufCommon(*this, fpTy, "fp", + /*allowLowPrecision=*/isA5))) + return failure(); auto fpSpace = getSpace(fpTy); if (!fpSpace || *fpSpace != pto::AddressSpace::SCALING) return emitOpError("expects fp to be loc=scaling"); @@ -6879,10 +6918,10 @@ mlir::LogicalResult mlir::pto::TInsertOp::verify() { if (dstTb.getSFractalSizeI32() != 512) return emitOpError("expects A2/A3 tinsert dst fractal size to be 512"); - if (hasPreQuantScalar) { + if (hasFp || hasPreQuantScalar) { if (!isA2A3AccQuantInsertTypePair(srcElem, dstElem)) return emitOpError( - "expects A2/A3 acc preQuantScalar tinsert element types to be " + "expects A2/A3 acc fp/preQuantScalar tinsert element types to be " "(src=f32,dst=i8) or (src=i32,dst=i8/f16/i16)"); } else if (!isA2A3AccCastInsertTypePair(srcElem, dstElem)) { return emitOpError( @@ -7006,19 +7045,6 @@ static bool isColMajorRowMajorNZTileBuf(pto::TileBufType ty) { ty.getSLayoutValueI32() == static_cast(pto::SLayout::RowMajor); } -static bool isRowMajorNoneBoxNDTileBuf(pto::TileBufType ty) { - return ty.getBLayoutValueI32() == static_cast(pto::BLayout::RowMajor) && - ty.getSLayoutValueI32() == static_cast(pto::SLayout::NoneBox); -} - -static bool isA2A3VectorPreQuantTypePair(Type srcElem, Type dstElem) { - if (srcElem.isF32()) - return dstElem.isInteger(8); - if (srcElem.isInteger(32)) - return dstElem.isInteger(8) || dstElem.isF16() || dstElem.isInteger(16); - return false; -} - static bool isA5Fp8LikeType(Type ty) { if (auto ft = dyn_cast(ty)) return ft.getWidth() == 8; @@ -7063,212 +7089,6 @@ static bool isA5VectorPreQuantTypePair(Type srcElem, Type dstElem) { return false; } -mlir::LogicalResult mlir::pto::TExtractFPOp::verify() { - auto verifyCommon = [&](bool allowLowPrecision) - -> FailureOr> { - Type srcTy = getSrc().getType(); - Type fpTy = getFp().getType(); - Type dstTy = getDst().getType(); - auto srcTb = dyn_cast(srcTy); - auto fpTb = dyn_cast(fpTy); - auto dstTb = dyn_cast(dstTy); - if (!srcTb || !fpTb || !dstTb) - return emitOpError("expects src, fp, and dst to be !pto.tile_buf"); - if (failed(verifyTileBufCommon(*this, srcTy, "src", allowLowPrecision)) || - failed(verifyTileBufCommon(*this, fpTy, "fp", allowLowPrecision)) || - failed(verifyTileBufCommon(*this, dstTy, "dst", allowLowPrecision)) || - failed(verifyNonNegativeIndexRowCol( - *getOperation(), getIndexRow(), getIndexCol(), - /*includeIndexAndIntOpsInConstFold=*/true)) || - failed(verifyExtractStaticBoundsCommon( - *getOperation(), getIndexRow(), getIndexCol(), srcTy, dstTy, - /*includeIndexAndIntOpsInConstFold=*/true))) - return failure(); - auto srcSpace = getPTOMemorySpaceEnum(srcTy); - auto fpSpace = getPTOMemorySpaceEnum(fpTy); - auto dstSpace = getPTOMemorySpaceEnum(dstTy); - if (!srcSpace || !fpSpace || !dstSpace) - return emitOpError("expects src, fp, and dst to have explicit loc"); - if (*srcSpace != pto::AddressSpace::ACC) - return emitOpError("expects src to use loc=acc"); - if (*fpSpace != pto::AddressSpace::SCALING) - return emitOpError("expects fp to use loc=scaling"); - if (*dstSpace != pto::AddressSpace::MAT && *dstSpace != pto::AddressSpace::VEC) - return emitOpError("expects dst to use loc=mat or loc=vec"); - if (!isColMajorRowMajorNZTileBuf(srcTb)) - return emitOpError("expects src to use blayout=col_major and slayout=row_major"); - if (*dstSpace == pto::AddressSpace::MAT) { - if (!isColMajorRowMajorNZTileBuf(dstTb)) - return emitOpError("expects mat dst to use blayout=col_major and slayout=row_major"); - } else { - if (!(dstTb.getBLayoutValueI32() == static_cast(pto::BLayout::RowMajor) && - dstTb.getSLayoutValueI32() == static_cast(pto::SLayout::NoneBox))) - return emitOpError("expects vec dst to use ND layout (blayout=row_major, slayout=none_box)"); - } - return std::make_tuple(srcTy, fpTy, dstTy, srcTb, fpTb, dstTb, *srcSpace, - *fpSpace, *dstSpace); - }; - auto accToVecModeAttr = getAccToVecModeAttr(); - const bool hasAccToVecMode = static_cast(accToVecModeAttr); - auto verifyA2A3 = [&]() -> LogicalResult { - auto common = verifyCommon(/*allowLowPrecision=*/false); - if (failed(common)) - return failure(); - auto [srcTy, fpTy, dstTy, srcTb, fpTb, dstTb, srcSpace, fpSpace, dstSpace] = - *common; - (void)fpTy; - (void)srcSpace; - (void)fpSpace; - (void)dstSpace; - if (hasAccToVecMode) - return emitOpError("expects accToVecMode only on A5 acc->vec textract_fp forms"); - if (dstSpace != pto::AddressSpace::MAT) - return emitOpError("expects A2/A3 textract_fp dst to use loc=mat"); - if (dstTb.getSFractalSizeI32() != 512) - return emitOpError("expects dst fractal size to be 512"); - if (hasAccToVecMode && dstSpace != pto::AddressSpace::VEC) - return emitOpError("expects accToVecMode only on A5 acc->vec textract_fp forms"); - Type srcElem = getElemTy(srcTy); - Type dstElem = getElemTy(dstTy); - if (!isA2A3VectorPreQuantTypePair(srcElem, dstElem)) - return emitOpError( - "expects A2/A3 textract_fp element types to be (src=f32,dst=i8) " - "or (src=i32,dst=i8/f16/i16)"); - return success(); - }; - auto verifyA5 = [&]() -> LogicalResult { - auto common = verifyCommon(/*allowLowPrecision=*/true); - if (failed(common)) - return failure(); - auto [srcTy, fpTy, dstTy, srcTb, fpTb, dstTb, srcSpace, fpSpace, dstSpace] = - *common; - (void)fpTy; - (void)srcTb; - (void)fpTb; - (void)dstTb; - (void)srcSpace; - (void)fpSpace; - (void)dstSpace; - Type srcElem = getElemTy(srcTy); - Type dstElem = getElemTy(dstTy); - if (!isA5VectorPreQuantTypePair(srcElem, dstElem)) - return emitOpError( - "expects A5 textract_fp element types to be (src=f32,dst=i8/fp8/f16/bf16/f32) " - "or (src=i32,dst=i8/f16/bf16)"); - return success(); - }; - return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); -} - -mlir::LogicalResult mlir::pto::TInsertFPOp::verify() { - auto verifyCommon = [&](bool allowLowPrecision, bool isA5) - -> FailureOr> { - Type srcTy = getSrc().getType(); - Type fpTy = getFp().getType(); - Type dstTy = getDst().getType(); - auto srcTb = dyn_cast(srcTy); - auto fpTb = dyn_cast(fpTy); - auto dstTb = dyn_cast(dstTy); - if (!srcTb || !fpTb || !dstTb) - return emitOpError("expects src, fp, and dst to be !pto.tile_buf"); - if (failed(verifyTileBufCommon(*this, srcTy, "src", allowLowPrecision)) || - failed(verifyTileBufCommon(*this, fpTy, "fp", allowLowPrecision)) || - failed(verifyTileBufCommon(*this, dstTy, "dst", allowLowPrecision)) || - failed(verifyNonNegativeIndexRowCol( - *getOperation(), getIndexRow(), getIndexCol(), - /*includeIndexAndIntOpsInConstFold=*/true)) || - failed(verifyInsertStaticBoundsCommon( - *getOperation(), getIndexRow(), getIndexCol(), srcTy, dstTy, - /*includeIndexAndIntOpsInConstFold=*/true))) - return failure(); - auto srcSpace = getPTOMemorySpaceEnum(srcTy); - auto fpSpace = getPTOMemorySpaceEnum(fpTy); - auto dstSpace = getPTOMemorySpaceEnum(dstTy); - if (!srcSpace || !fpSpace || !dstSpace) - return emitOpError("expects src, fp, and dst to have explicit loc"); - if (*srcSpace != pto::AddressSpace::ACC) - return emitOpError("expects src to use loc=acc"); - if (*fpSpace != pto::AddressSpace::SCALING) - return emitOpError("expects fp to use loc=scaling"); - // A2/A3: only acc->mat; A5: acc->mat or acc->vec. - if (*dstSpace != pto::AddressSpace::MAT && - !(isA5 && *dstSpace == pto::AddressSpace::VEC)) - return emitOpError("expects dst to use loc=mat" + - (isA5 ? StringRef(" or loc=vec (A5)") : StringRef(""))); - if (!isColMajorRowMajorNZTileBuf(srcTb)) - return emitOpError("expects src to use blayout=col_major and slayout=row_major"); - if (*dstSpace == pto::AddressSpace::MAT && !isColMajorRowMajorNZTileBuf(dstTb)) - return emitOpError("expects dst (mat) to use blayout=col_major and slayout=row_major"); - if (*dstSpace == pto::AddressSpace::VEC && - !isRowMajorNoneBoxNDTileBuf(dstTb) && !isColMajorRowMajorNZTileBuf(dstTb)) - return emitOpError("expects dst (vec) to use ND(row_major/none_box) or NZ(col_major/row_major) layout"); - // accToVecMode is only valid when dst=vec. - if (static_cast(getAccToVecModeAttr()) && - *dstSpace != pto::AddressSpace::VEC) - return emitOpError("accToVecMode is only valid with dst=vec"); - return std::make_tuple(srcTy, fpTy, dstTy, srcTb, fpTb, dstTb, *srcSpace, - *fpSpace, *dstSpace); - }; - auto accToVecModeAttr = getAccToVecModeAttr(); - const bool hasAccToVecMode = static_cast(accToVecModeAttr); - auto verifyA2A3 = [&]() -> LogicalResult { - auto common = verifyCommon(/*allowLowPrecision=*/false, /*isA5=*/false); - if (failed(common)) - return failure(); - auto [srcTy, fpTy, dstTy, srcTb, fpTb, dstTb, srcSpace, fpSpace, dstSpace] = - *common; - (void)fpTy; - (void)srcTb; - (void)fpTb; - (void)srcSpace; - (void)fpSpace; - (void)dstSpace; - if (hasAccToVecMode) - return emitOpError("expects accToVecMode only on A5 acc->vec tinsert_fp forms"); - if (dstSpace != pto::AddressSpace::MAT) - return emitOpError("expects A2/A3 tinsert_fp dst to use loc=mat"); - if (dstTb.getSFractalSizeI32() != 512) - return emitOpError("expects dst fractal size to be 512"); - if (hasAccToVecMode && dstSpace != pto::AddressSpace::VEC) - return emitOpError("expects accToVecMode only on A5 acc->vec tinsert_fp forms"); - Type srcElem = getElemTy(srcTy); - Type dstElem = getElemTy(dstTy); - if (!isA2A3VectorPreQuantTypePair(srcElem, dstElem)) - return emitOpError( - "expects A2/A3 tinsert_fp element types to be (src=f32,dst=i8) " - "or (src=i32,dst=i8/f16/i16)"); - return success(); - }; - auto verifyA5 = [&]() -> LogicalResult { - auto common = verifyCommon(/*allowLowPrecision=*/true, /*isA5=*/true); - if (failed(common)) - return failure(); - auto [srcTy, fpTy, dstTy, srcTb, fpTb, dstTb, srcSpace, fpSpace, dstSpace] = - *common; - (void)fpTy; - (void)srcTb; - (void)fpTb; - (void)dstTb; - (void)srcSpace; - (void)fpSpace; - (void)dstSpace; - Type srcElem = getElemTy(srcTy); - Type dstElem = getElemTy(dstTy); - if (!isA5VectorPreQuantTypePair(srcElem, dstElem)) - return emitOpError( - "expects A5 tinsert_fp element types to be (src=f32,dst=i8/fp8/f16/bf16/f32) " - "or (src=i32,dst=i8/f16/bf16)"); - return success(); - }; - return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); -} - static mlir::LogicalResult verifyTFillPadLike(Operation *op, Type srcTy, Type dstTy, pto::TFillPadMode mode) { @@ -7841,6 +7661,13 @@ mlir::LogicalResult mlir::pto::TMovOp::verify() { return emitOpError() << "expects acc-source fp/relu tmov src to use blayout=col_major and slayout=row_major"; } + if (hasFp && !isA5 && dstTb && isAccToMat && + (dstTb.getBLayoutValueI32() != + static_cast(pto::BLayout::ColMajor) || + dstTb.getSLayoutValueI32() != + static_cast(pto::SLayout::RowMajor))) + return emitOpError() + << "expects fp tmov dst to use blayout=col_major and slayout=row_major"; if (srcTb && dstTb && isAccToMat && !isA5 && dstTb.getSFractalSizeI32() != 512) return emitOpError() << "expects A2/A3 acc-to-mat tmov destination fractal to be 512"; @@ -7852,73 +7679,6 @@ mlir::LogicalResult mlir::pto::TMovOp::verify() { return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); } -mlir::LogicalResult mlir::pto::TMovFPOp::verify() { - auto verifyA2A3 = [&]() -> LogicalResult { - Type srcTy = getSrc().getType(); - Type fpTy = getFp().getType(); - Type dstTy = getDst().getType(); - if (failed(verifyTileBufCommon(*this, srcTy, "src")) || - failed(verifyTileBufCommon(*this, fpTy, "fp")) || - failed(verifyTileBufCommon(*this, dstTy, "dst"))) - return failure(); - auto srcElemTy = getElemTy(srcTy); - auto srcIntTy = dyn_cast(srcElemTy); - if (!(srcElemTy.isF32() || - (srcIntTy && srcIntTy.getWidth() == 32))) - return emitOpError() - << "expects src to have element type f32, i32"; - auto fpSpace = getPTOMemorySpaceEnum(fpTy); - if (!fpSpace || *fpSpace != mlir::pto::AddressSpace::SCALING) - return emitOpError() << "expects fp to be in the scaling address space"; - auto srcSpace = getPTOMemorySpaceEnum(srcTy); - if (!srcSpace || *srcSpace != mlir::pto::AddressSpace::ACC) - return emitOpError() << "expects src to be in the acc address space"; - auto dstSpace = getPTOMemorySpaceEnum(dstTy); - if (!dstSpace || *dstSpace != mlir::pto::AddressSpace::MAT) - return emitOpError() << "expects dst to be in the mat address space"; - auto srcTb = dyn_cast(srcTy); - auto dstTb = dyn_cast(dstTy); - if (srcTb && - (srcTb.getBLayoutValueI32() != static_cast(pto::BLayout::ColMajor) || - srcTb.getSLayoutValueI32() != static_cast(pto::SLayout::RowMajor))) - return emitOpError() - << "expects src to use blayout=col_major and slayout=row_major"; - if (dstTb && - (dstTb.getBLayoutValueI32() != static_cast(pto::BLayout::ColMajor) || - dstTb.getSLayoutValueI32() != static_cast(pto::SLayout::RowMajor))) - return emitOpError() - << "expects dst to use blayout=col_major and slayout=row_major"; - if (dstTb && dstTb.getSFractalSizeI32() != 512) - return emitOpError() << "expects dst to use fractal size 512"; - return mlir::success(); - }; - auto verifyA5 = [&]() -> LogicalResult { - Type srcTy = getSrc().getType(); - Type fpTy = getFp().getType(); - Type dstTy = getDst().getType(); - if (failed(verifyTileBufCommon(*this, srcTy, "src", /*allowLowPrecision=*/true)) || - failed(verifyTileBufCommon(*this, fpTy, "fp", /*allowLowPrecision=*/true)) || - failed(verifyTileBufCommon(*this, dstTy, "dst", /*allowLowPrecision=*/true))) - return failure(); - auto srcElemTy = getElemTy(srcTy); - auto srcIntTy = dyn_cast(srcElemTy); - if (!(srcElemTy.isF32() || - (srcIntTy && srcIntTy.getWidth() == 32))) - return emitOpError() - << "expects src to have element type f32, i32"; - auto fpSpace = getPTOMemorySpaceEnum(fpTy); - if (!fpSpace || *fpSpace != mlir::pto::AddressSpace::SCALING) - return emitOpError() << "expects fp to be in the scaling address space"; - auto srcTb = dyn_cast(srcTy); - if (srcTb && - (srcTb.getBLayoutValueI32() != static_cast(pto::BLayout::ColMajor) || - srcTb.getSLayoutValueI32() != static_cast(pto::SLayout::RowMajor))) - return emitOpError() - << "expects src to use blayout=col_major and slayout=row_major"; - return success(); - }; - return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); -} // 辅助函数：获取 Rank，支持 ShapedType 和 PTO TileTypes static int64_t getRankHelper(Type t) { if (auto s = dyn_cast(t)) return s.getRank(); @@ -12357,103 +12117,6 @@ mlir::LogicalResult mlir::pto::TSqrtOp::verify() { return mlir::success(); } -mlir::LogicalResult mlir::pto::TStoreFPOp::verify() { - auto verifySrcDtypeAlways = [&]() -> LogicalResult { - Type srcTy = getSrc().getType(); - auto srcElemTy = getElemTy(srcTy); - if (!srcElemTy) - return success(); - auto srcIntTy = dyn_cast(srcElemTy); - if (!(srcElemTy.isF32() || - (srcIntTy && srcIntTy.getWidth() == 32))) - return emitOpError() - << "expects src to have element type f32, i32"; - return success(); - }; - - if (failed(verifySrcDtypeAlways())) - return failure(); - - auto verifyDstType = [&]() -> LogicalResult { - auto dstPart = dyn_cast(getDst().getType()); - if (!dstPart) - return emitOpError() << "expects dst to be !pto.partition_tensor_view"; - for (auto [idx, dim] : llvm::enumerate(dstPart.getShape())) { - if (dim != ShapedType::kDynamic && dim <= 0) - return emitOpError() - << "expects dst shape[" << idx << "] to be positive"; - } - return success(); - }; - - auto verifyA2A3 = [&]() -> LogicalResult { - Type srcTy = getSrc().getType(); - Type fpTy = getFp().getType(); - if (!isa(srcTy)) - return emitOpError() << "expects src to be a !pto.tile_buf"; - if (!isa(fpTy)) - return emitOpError() << "expects fp to be a !pto.tile_buf"; - if (failed(verifyTileBufCommon(*this, srcTy, "src")) || - failed(verifyTileBufCommon(*this, fpTy, "fp"))) - return failure(); - if (failed(verifyDstType())) - return failure(); - auto srcSpace = getPTOMemorySpaceEnum(srcTy); - if (!srcSpace || *srcSpace != pto::AddressSpace::ACC) - return emitOpError() << "expects src to be in the acc address space"; - auto srcElemTy = getElemTy(srcTy); - auto srcIntTy = dyn_cast(srcElemTy); - if (!(srcElemTy.isF32() || - (srcIntTy && srcIntTy.getWidth() == 32))) - return emitOpError() - << "expects src to have element type f32, i32"; - auto srcShape = getShapeVec(srcTy); - if (srcShape.size() != 2) - return emitOpError() << "expects src to have rank 2"; - if (srcShape[1] != ShapedType::kDynamic && - (srcShape[1] < 1 || srcShape[1] > 4095)) - return emitOpError() << "expects src.cols to be in the range [1, 4095]"; - auto srcValid = getValidShapeVec(srcTy); - if (srcValid.size() != 2) - return emitOpError() << "expects src to have a rank-2 valid_shape"; - if (srcValid[1] != ShapedType::kDynamic && - (srcValid[1] < 0 || srcValid[1] > 4095)) - return emitOpError() - << "expects src.valid_shape[1] to be in the range [0, 4095]"; - return mlir::success(); - }; - auto verifyA5 = [&]() -> LogicalResult { - Type srcTy = getSrc().getType(); - Type fpTy = getFp().getType(); - if (!isa(srcTy)) - return emitOpError() << "expects src to be a !pto.tile_buf"; - if (!isa(fpTy)) - return emitOpError() << "expects fp to be a !pto.tile_buf"; - if (failed(verifyTileBufCommon(*this, srcTy, "src")) || - failed(verifyTileBufCommon(*this, fpTy, "fp"))) - return failure(); - if (failed(verifyDstType())) - return failure(); - auto srcSpace = getPTOMemorySpaceEnum(srcTy); - if (!srcSpace || *srcSpace != pto::AddressSpace::ACC) - return emitOpError() << "expects src to be in the acc address space"; - auto srcElemTy = getElemTy(srcTy); - auto srcIntTy = dyn_cast(srcElemTy); - if (!(srcElemTy.isF32() || - (srcIntTy && srcIntTy.getWidth() == 32))) - return emitOpError() - << "expects src to have element type f32, i32"; - return mlir::success(); - }; - switch (getVerifierTargetArch(getOperation())) { - case VerifierTargetArch::A2A3: - return verifyA2A3(); - case VerifierTargetArch::A5: - return verifyA5(); - } - return failure(); -} - mlir::LogicalResult mlir::pto::TSubOp::verify() { return verifyArithmeticBinaryTileOpWithArchDispatch( getOperation(), getSrc0().getType(), getSrc1().getType(), getDst().getType(), @@ -13862,21 +13525,15 @@ void TAbsOp::getEffects( // Read: src, Write: dst (GM) void TStoreOp::getEffects(SmallVectorImpl> &effects) { addEffect(effects, &getSrcMutable(), MemoryEffects::Read::get()); + auto fpRange = getFpMutable(); + if (!fpRange.empty()) + addEffect(effects, &*fpRange.begin(), MemoryEffects::Read::get()); auto preQuantRange = getPreQuantScalarMutable(); if (!preQuantRange.empty()) addEffect(effects, &*preQuantRange.begin(), MemoryEffects::Read::get()); addEffect(effects, &getDstMutable(), MemoryEffects::Write::get()); } -// === TStoreFPOp === -// Read: src/fp, Write: dst (GM) -void TStoreFPOp::getEffects( - SmallVectorImpl> &effects) { - addEffect(effects, &getSrcMutable(), MemoryEffects::Read::get()); - addEffect(effects, &getFpMutable(), MemoryEffects::Read::get()); - addEffect(effects, &getDstMutable(), MemoryEffects::Write::get()); -} - // === TMovOp === // Read: src, Write: dst void TMovOp::getEffects(SmallVectorImpl> &effects) { @@ -14098,8 +13755,11 @@ void TExpandsOp::getEffects( // TEXTRACT: Read(src) -> Write(dst) void TExtractOp::getEffects( SmallVectorImpl> &effects) { - PTO_ADD_READ(getSrcMutable()); - PTO_ADD_WRITE(getDstMutable()); + addEffect(effects, &getSrcMutable(), MemoryEffects::Read::get()); + auto fpRange = getFpMutable(); + if (!fpRange.empty()) + addEffect(effects, &*fpRange.begin(), MemoryEffects::Read::get()); + addEffect(effects, &getDstMutable(), MemoryEffects::Write::get()); } // TINSERT: Read(src) -> Write(dst) @@ -14112,22 +13772,6 @@ void TInsertOp::getEffects( addEffect(effects, &getDstMutable(), MemoryEffects::Write::get()); } -// TEXTRACT_FP: Read(src), Read(fp) -> Write(dst) -void TExtractFPOp::getEffects( - SmallVectorImpl> &effects) { - PTO_ADD_READ(getSrcMutable()); - PTO_ADD_READ(getFpMutable()); - PTO_ADD_WRITE(getDstMutable()); -} - -// TINSERT_FP: Read(src), Read(fp) -> Write(dst) -void TInsertFPOp::getEffects( - SmallVectorImpl> &effects) { - PTO_ADD_READ(getSrcMutable()); - PTO_ADD_READ(getFpMutable()); - PTO_ADD_WRITE(getDstMutable()); -} - PTO_DEFINE_UNARY_EFFECTS(TFillPadOp, getSrcMutable(), getDstMutable()) void TGatherOp::getEffects( @@ -14151,8 +13795,6 @@ PTO_DEFINE_UNARY_EFFECTS(TMaxSOp, getSrcMutable(), getDstMutable()) PTO_DEFINE_BINARY_EFFECTS(TMinOp, getSrc0Mutable(), getSrc1Mutable(), getDstMutable()) PTO_DEFINE_UNARY_EFFECTS(TMinSOp, getSrcMutable(), getDstMutable()) -PTO_DEFINE_BINARY_EFFECTS(TMovFPOp, getSrcMutable(), getFpMutable(), getDstMutable()) - void TMrgSortOp::getEffects( SmallVectorImpl> &effects) { for (auto &opnd : getSrcsMutable()) { diff --git a/lib/PTO/Transforms/ConvertToPTOOp.cpp b/lib/PTO/Transforms/ConvertToPTOOp.cpp index 7e72901f56..9fb56eef67 100644 --- a/lib/PTO/Transforms/ConvertToPTOOp.cpp +++ b/lib/PTO/Transforms/ConvertToPTOOp.cpp @@ -149,7 +149,7 @@ struct MemrefCopyOpLowering : public OpRewritePattern { bool convertToStore = isFromFunctionArg(dst); if (convertToStore) { rewriter.replaceOpWithNewOp(copyOp, TypeRange(), src, dst, - Value{}); + Value{}, Value{}); return success(); } @@ -174,7 +174,7 @@ struct BufferizeMaterializeOpLowering if (convertToStore) { rewriter.replaceOpWithNewOp(bufMIDOp, TypeRange(), bufMIDOp.getSource(), dst, - Value{}); + Value{}, Value{}); return success(); } return failure(); diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 2d5bd758e1..a2a4e0105e 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -4933,6 +4933,9 @@ struct PTOTStoreToTSTORE : public OpConversionPattern { auto *ctx = rewriter.getContext(); Value src = peelUnrealized(adaptor.getSrc()); Value dst = peelUnrealized(adaptor.getDst()); + Value fp; + if (op.getFp()) + fp = peelUnrealized(adaptor.getFp()); Value preQuantScalar; if (op.getPreQuantScalar()) preQuantScalar = peelUnrealized(adaptor.getPreQuantScalar()); @@ -4941,6 +4944,7 @@ struct PTOTStoreToTSTORE : public OpConversionPattern { const auto phase = op.getStPhase(); const auto atomicType = op.getAtomicType(); const auto reluPreMode = op.getReluPreMode(); + const bool hasFp = static_cast(fp); const bool hasPreQuantScalar = static_cast(preQuantScalar); const bool phaseNonDefault = phase != pto::STPhase::Unspecified; const bool atomicNonDefault = atomicType != pto::AtomicType::AtomicNone; @@ -4953,6 +4957,34 @@ struct PTOTStoreToTSTORE : public OpConversionPattern { }; ArrayAttr targs; + if (hasFp) { + SmallVector operands{dstArg, src, fp}; + if (atomicNonDefault || reluNonDefault) { + auto srcTokOr = getOpaqueTok(src, "src"); + auto dstTokOr = getOpaqueTok(dstArg, "dst"); + auto fpTokOr = getOpaqueTok(fp, "fp"); + if (failed(srcTokOr) || failed(dstTokOr) || failed(fpTokOr)) + return failure(); + targs = rewriter.getArrayAttr({ + emitc::OpaqueAttr::get(ctx, *srcTokOr), + emitc::OpaqueAttr::get(ctx, *dstTokOr), + emitc::OpaqueAttr::get(ctx, *fpTokOr), + emitc::OpaqueAttr::get(ctx, atomicTypeTok(atomicType)), + emitc::OpaqueAttr::get(ctx, reluPreModeTok(reluPreMode)), + }); + } else { + targs = ArrayAttr{}; + } + + rewriter.create( + loc, TypeRange{}, "TSTORE_FP", ArrayAttr{}, targs, operands); + if (op->getNumResults() == 1) + rewriter.replaceOp(op, dst); + else + rewriter.eraseOp(op); + return success(); + } + // Map op attributes/operands to the exact TSTORE overload family: // 1) TSTORE(dst, src) // 2) TSTORE(dst, src) @@ -9621,14 +9653,20 @@ struct PTOExtractToEmitC : public OpConversionPattern { Value preQuantScalar; if (op.getPreQuantScalar()) preQuantScalar = peelUnrealized(adaptor.getPreQuantScalar()); + Value fp; + if (op.getFp()) + fp = peelUnrealized(adaptor.getFp()); auto modeAttr = op.getAccToVecModeAttr(); + const bool hasFp = static_cast(fp); const bool hasPreQuantScalar = static_cast(preQuantScalar); const bool hasMode = static_cast(modeAttr); const bool reluNonDefault = op.getReluPreMode() != pto::ReluPreMode::NoRelu; SmallVector operands{dst, src}; + if (hasFp) + operands.push_back(fp); if (hasPreQuantScalar) operands.push_back(preQuantScalar); operands.push_back(r0); @@ -9645,6 +9683,13 @@ struct PTOExtractToEmitC : public OpConversionPattern { emitc::OpaqueAttr::get(ctx, dstOT.getValue().str()), emitc::OpaqueAttr::get(ctx, srcOT.getValue().str()), }; + if (hasFp) { + auto fpOT = mlir::dyn_cast(fp.getType()); + if (!fpOT) + return rewriter.notifyMatchFailure( + op, "textract template lowering expects opaque fp type"); + args.push_back(emitc::OpaqueAttr::get(ctx, fpOT.getValue().str())); + } if (hasMode) args.push_back(emitc::OpaqueAttr::get(ctx, getAccToVecModeToken(modeAttr.getValue()))); args.push_back(emitc::OpaqueAttr::get(ctx, getReluPreModeToken(op.getReluPreMode()))); @@ -9652,53 +9697,8 @@ struct PTOExtractToEmitC : public OpConversionPattern { } rewriter.create( - loc, TypeRange{}, "TEXTRACT", ArrayAttr{}, templateArgs, operands); - rewriter.eraseOp(op); - return success(); - } -}; - -struct PTOExtractFPToEmitC : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult matchAndRewrite(pto::TExtractFPOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto loc = op.getLoc(); - auto *ctx = rewriter.getContext(); - - Value src = peelUnrealized(adaptor.getSrc()); - Value fp = peelUnrealized(adaptor.getFp()); - Value dst = peelUnrealized(adaptor.getDst()); - Value r0 = peelUnrealized(adaptor.getIndexRow()); - Value c0 = peelUnrealized(adaptor.getIndexCol()); - - auto modeAttr = op.getAccToVecModeAttr(); - const bool hasMode = static_cast(modeAttr); - const bool reluNonDefault = - op.getReluPreMode() != pto::ReluPreMode::NoRelu; - - ArrayAttr templateArgs; - if (hasMode || reluNonDefault) { - auto dstOT = mlir::dyn_cast(dst.getType()); - auto srcOT = mlir::dyn_cast(src.getType()); - auto fpOT = mlir::dyn_cast(fp.getType()); - if (!dstOT || !srcOT || !fpOT) - return rewriter.notifyMatchFailure( - op, "textract_fp template lowering expects opaque dst/src/fp types"); - SmallVector args{ - emitc::OpaqueAttr::get(ctx, dstOT.getValue().str()), - emitc::OpaqueAttr::get(ctx, srcOT.getValue().str()), - emitc::OpaqueAttr::get(ctx, fpOT.getValue().str()), - }; - if (hasMode) - args.push_back(emitc::OpaqueAttr::get(ctx, getAccToVecModeToken(modeAttr.getValue()))); - args.push_back(emitc::OpaqueAttr::get(ctx, getReluPreModeToken(op.getReluPreMode()))); - templateArgs = rewriter.getArrayAttr(args); - } - - rewriter.create( - loc, TypeRange{}, hasMode ? "TEXTRACT" : "TEXTRACT_FP", ArrayAttr{}, templateArgs, - ValueRange{dst, src, fp, r0, c0}); + loc, TypeRange{}, hasFp && !hasMode ? "TEXTRACT_FP" : "TEXTRACT", + ArrayAttr{}, templateArgs, operands); rewriter.eraseOp(op); return success(); } @@ -9767,63 +9767,17 @@ struct PTOInsertToEmitC : public OpConversionPattern { templateArgs = rewriter.getArrayAttr(args); } - rewriter.create( - loc, TypeRange{}, "TINSERT", ArrayAttr{}, templateArgs, operands); - rewriter.eraseOp(op); - return success(); - } -}; - -struct PTOInsertFPToEmitC : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult matchAndRewrite(pto::TInsertFPOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto loc = op.getLoc(); - auto *ctx = rewriter.getContext(); - - Value src = peelUnrealized(adaptor.getSrc()); - Value fp = peelUnrealized(adaptor.getFp()); - Value dst = peelUnrealized(adaptor.getDst()); - Value r0 = peelUnrealized(adaptor.getIndexRow()); - Value c0 = peelUnrealized(adaptor.getIndexCol()); - - auto modeAttr = op.getAccToVecModeAttr(); - const bool hasMode = static_cast(modeAttr); - const bool reluNonDefault = - op.getReluPreMode() != pto::ReluPreMode::NoRelu; - - ArrayAttr templateArgs = ArrayAttr{}; - if (hasMode || reluNonDefault) { - auto dstOT = mlir::dyn_cast(dst.getType()); - auto srcOT = mlir::dyn_cast(src.getType()); - auto fpOT = mlir::dyn_cast(fp.getType()); - if (!dstOT || !srcOT || !fpOT) - return rewriter.notifyMatchFailure( - op, "tinsert_fp template lowering expects opaque dst/src/fp types"); - SmallVector args{ - emitc::OpaqueAttr::get(ctx, dstOT.getValue().str()), - emitc::OpaqueAttr::get(ctx, srcOT.getValue().str()), - emitc::OpaqueAttr::get(ctx, fpOT.getValue().str()), - }; - if (hasMode) - args.push_back(emitc::OpaqueAttr::get(ctx, getAccToVecModeToken(modeAttr.getValue()))); - args.push_back(emitc::OpaqueAttr::get(ctx, getReluPreModeToken(op.getReluPreMode()))); - templateArgs = rewriter.getArrayAttr(args); - } + if (hasFp && !hasMode && !reluNonDefault) + templateArgs = ArrayAttr{}; rewriter.create( - loc, TypeRange{}, hasMode ? "TINSERT" : "TINSERT_FP", ArrayAttr{}, templateArgs, - ValueRange{dst, src, fp, r0, c0}); + loc, TypeRange{}, hasFp && !hasMode ? "TINSERT_FP" : "TINSERT", + ArrayAttr{}, templateArgs, operands); rewriter.eraseOp(op); return success(); } }; -//===----------------------------------------------------------------------===// -// pto.tfillpad lowering -> TFILLPAD(dst, src) -//===----------------------------------------------------------------------===// - static StringRef getTFillPadModeToken(pto::TFillPadMode mode) { switch (mode) { case pto::TFillPadMode::Normal: @@ -10301,44 +10255,6 @@ struct PTOMovToEmitC : public OpConversionPattern { // PTOConvert.cpp (add lowering + patterns.add for TMOV_FP DPS/memref op) //===----------------------------------------------------------------------===// -struct PTOMovFPToEmitC : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult matchAndRewrite(pto::TMovFPOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto loc = op.getLoc(); - auto *ctx = rewriter.getContext(); - - Value dst = peelUnrealized(adaptor.getDst()); - Value src = peelUnrealized(adaptor.getSrc()); - Value fp = peelUnrealized(adaptor.getFp()); - - // TMOV_FP(dstTileData, cTile, fbTile) - ArrayAttr templateArgs; - auto dstOT = mlir::dyn_cast(dst.getType()); - auto srcOT = mlir::dyn_cast(src.getType()); - auto fpOT = mlir::dyn_cast(fp.getType()); - if (dstOT && srcOT && fpOT) { - templateArgs = rewriter.getArrayAttr({ - emitc::OpaqueAttr::get(ctx, dstOT.getValue().str()), - emitc::OpaqueAttr::get(ctx, srcOT.getValue().str()), - emitc::OpaqueAttr::get(ctx, fpOT.getValue().str()), - }); - } else { - templateArgs = ArrayAttr{}; - } - - SmallVector operands{dst, src, fp}; - rewriter.create( - loc, TypeRange{}, "TMOV_FP", - /*args=*/ArrayAttr{}, /*templateArgs=*/templateArgs, - /*operands=*/operands); - - rewriter.eraseOp(op); - return success(); - } -}; - struct PTOQuantToEmitC : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -12070,32 +11986,6 @@ struct PTOSqrtSToEmitC : public OpConversionPattern { // PTOConvert.cpp (add lowering + patterns.add for TSTORE_FP DPS/memref op) //===----------------------------------------------------------------------===// -struct PTOStoreFPSToEmitC : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult matchAndRewrite(pto::TStoreFPOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto loc = op.getLoc(); - - Value src = peelUnrealized(adaptor.getSrc()); - Value fp = peelUnrealized(adaptor.getFp()); - Value dst = peelUnrealized(adaptor.getDst()); - - SmallVector operands{dst, src, fp}; - rewriter.create( - loc, TypeRange{}, "TSTORE_FP", - /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, - /*operands=*/operands); - - rewriter.eraseOp(op); - return success(); - } -}; - -//===----------------------------------------------------------------------===// -// PTOConvert.cpp (add lowering + patterns.add for TSUB DPS/memref op) -//===----------------------------------------------------------------------===// - struct PTOSubSToEmitC : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -13369,7 +13259,6 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); - patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); @@ -13425,12 +13314,10 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); - patterns.add(typeConverter, ctx); + patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); - patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); diff --git a/lib/TileOps/__init__.py b/lib/TileOps/__init__.py index f595bfd22d..083986769e 100644 --- a/lib/TileOps/__init__.py +++ b/lib/TileOps/__init__.py @@ -41,8 +41,7 @@ ("a5", "pto.tdequant"): ".a5.tdequant", ("a5", "pto.texp"): ".a5.texp", ("a5", "pto.texpands"): ".a5.texpand", - ("a5", "pto.textract"): ".a5.textract", - ("a5", "pto.textract_fp"): ".a5.textract_fp", + ("a5", "pto.textract"): (".a5.textract", ".a5.textract_fp"), ("a5", "pto.tfmod"): ".a5.tfmod", ("a5", "pto.tfmods"): ".a5.tfmods", ("a5", "pto.tfillpad"): ".a5.tfillpad", @@ -118,7 +117,6 @@ ("a5", "pto.tmrgsort"): ".a5.tmrgsort", ("a5", "pto.tsort32"): ".a5.tsort32", ("a5", "pto.tstore"): ".a5.tstore", - ("a5", "pto.tstore_fp"): ".a5.tstore", ("a5", "pto.tsub"): ".a5.tsub", ("a5", "pto.tsubs"): ".a5.tsubs", ("a5", "pto.tsqrt"): ".a5.tsqrt", diff --git a/lib/TileOps/a5/textract_fp.py b/lib/TileOps/a5/textract_fp.py index eb67a7c1a2..b8280616a0 100644 --- a/lib/TileOps/a5/textract_fp.py +++ b/lib/TileOps/a5/textract_fp.py @@ -5,7 +5,7 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""PTODSL TileLib templates for ``pto.textract_fp``.""" +"""PTODSL TileLib templates for the ``pto.textract`` fp form.""" from ptodsl import pto import ptodsl.tilelib as tilelib @@ -27,7 +27,7 @@ def _register_textract_fp(name, signatures, quant_mode, template_id): signatures = (signatures,) @tilelib.tile_template( - op="pto.textract_fp", + op="pto.textract", target="a5", name=name, dtypes=tuple(signatures), @@ -40,7 +40,7 @@ def _register_textract_fp(name, signatures, quant_mode, template_id): is_post_update=False, tags=("extract", "acc", "mat", "fp"), ) - def _template(src: pto.Tile, fp: pto.Tile, index_row: pto.i32, index_col: pto.i32, dst: pto.Tile): + def _template(src: pto.Tile, index_row: pto.i32, index_col: pto.i32, dst: pto.Tile, fp: pto.Tile): m, n = dst.valid_shape src_ptr = src.as_ptr() if str(src.dtype) == "si32": @@ -60,39 +60,39 @@ def _template(src: pto.Tile, fp: pto.Tile, index_row: pto.i32, index_col: pto.i3 template_textract_fp_f32_si8 = _register_textract_fp( "template_textract_fp_f32_si8", - ("f32", "f32", "i32", "i32", "si8"), + ("f32", "i32", "i32", "si8", "f32"), "qf322b8_pre_vec", 0, ) template_textract_fp_f32_ui8 = _register_textract_fp( "template_textract_fp_f32_ui8", - ("f32", "f32", "i32", "i32", "ui8"), + ("f32", "i32", "i32", "ui8", "f32"), "qf322b8_pre_vec", 1, ) template_textract_fp_f32_f16 = _register_textract_fp( "template_textract_fp_f32_f16", - ("f32", "f32", "i32", "i32", "f16"), + ("f32", "i32", "i32", "f16", "f32"), "qf322f16_pre_vec", 2, ) template_textract_fp_f32_bf16 = _register_textract_fp( "template_textract_fp_f32_bf16", - ("f32", "f32", "i32", "i32", "bf16"), + ("f32", "i32", "i32", "bf16", "f32"), "qf322bf16_pre_vec", 3, ) template_textract_fp_f32_f32 = _register_textract_fp( "template_textract_fp_f32_f32", - ("f32", "f32", "i32", "i32", "f32"), + ("f32", "i32", "i32", "f32", "f32"), "qf322f32_pre_vec", 4, ) template_textract_fp_si32_si8 = _register_textract_fp( "template_textract_fp_si32_si8", ( - ("si32", "f32", "i32", "i32", "si8"), - ("i32", "f32", "i32", "i32", "si8"), + ("si32", "i32", "i32", "si8", "f32"), + ("i32", "i32", "i32", "si8", "f32"), ), "req8_vec", 5, @@ -100,8 +100,8 @@ def _template(src: pto.Tile, fp: pto.Tile, index_row: pto.i32, index_col: pto.i3 template_textract_fp_si32_ui8 = _register_textract_fp( "template_textract_fp_si32_ui8", ( - ("si32", "f32", "i32", "i32", "ui8"), - ("i32", "f32", "i32", "i32", "ui8"), + ("si32", "i32", "i32", "ui8", "f32"), + ("i32", "i32", "i32", "ui8", "f32"), ), "req8_vec", 6, @@ -109,8 +109,8 @@ def _template(src: pto.Tile, fp: pto.Tile, index_row: pto.i32, index_col: pto.i3 template_textract_fp_si32_f16 = _register_textract_fp( "template_textract_fp_si32_f16", ( - ("si32", "f32", "i32", "i32", "f16"), - ("i32", "f32", "i32", "i32", "f16"), + ("si32", "i32", "i32", "f16", "f32"), + ("i32", "i32", "i32", "f16", "f32"), ), "deqf16_vec", 7, @@ -118,8 +118,8 @@ def _template(src: pto.Tile, fp: pto.Tile, index_row: pto.i32, index_col: pto.i3 template_textract_fp_si32_bf16 = _register_textract_fp( "template_textract_fp_si32_bf16", ( - ("si32", "f32", "i32", "i32", "bf16"), - ("i32", "f32", "i32", "i32", "bf16"), + ("si32", "i32", "i32", "bf16", "f32"), + ("i32", "i32", "i32", "bf16", "f32"), ), "qs322bf16_pre_vec", 8, diff --git a/lib/TileOps/a5/tstore.py b/lib/TileOps/a5/tstore.py index 3b216d31dc..de1c91080d 100644 --- a/lib/TileOps/a5/tstore.py +++ b/lib/TileOps/a5/tstore.py @@ -5,7 +5,7 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""PTODSL TileLib templates for ``pto.tstore`` and ``pto.tstore_fp``.""" +"""PTODSL TileLib templates for ``pto.tstore``.""" from ptodsl import pto import ptodsl.tilelib as tilelib @@ -292,7 +292,7 @@ def template_tstore_acc_to_gm_nz2nz(src: pto.Tile, dst: pto.PartitionTensorView) @tilelib.tile_template( - op="pto.tstore_fp", + op="pto.tstore", target="a5", name="template_tstore_fp_acc_to_gm", dtypes=(("f32", "f16", "f16"), ("f32", "bf16", "bf16")), @@ -305,7 +305,7 @@ def template_tstore_acc_to_gm_nz2nz(src: pto.Tile, dst: pto.PartitionTensorView) is_post_update=False, tags=("store", "acc", "gm", "fp"), ) -def template_tstore_fp_acc_to_gm(src: pto.Tile, fp: pto.Tile, dst: pto.PartitionTensorView): +def template_tstore_fp_acc_to_gm(src: pto.Tile, dst: pto.PartitionTensorView, fp: pto.Tile): m, n = src.valid_shape strides = dst.strides quant_mode = "qf322bf16_pre_vec" if str(fp.dtype) == "bf16" else "qf322f16_pre_vec" diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 3e2d4593f3..460b24bf27 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -2931,14 +2931,21 @@ def tload(part, tile): _pto.TLoadOp(None, unwrap_surface_value(part), unwrap_surface_value(tile)) -def tstore(tile, part): +def tstore(tile, part, *, fp=None): """``pto.tstore ins(tile) outs(part)``.""" - _pto.TStoreOp(None, unwrap_surface_value(tile), unwrap_surface_value(part)) + kwargs = {} + if fp is not None: + kwargs["fp"] = unwrap_surface_value(fp) + _pto.TStoreOp( + None, unwrap_surface_value(tile), unwrap_surface_value(part), **kwargs + ) -def tmov(src, dst, *, mode=None): +def tmov(src, dst, *, fp=None, mode=None): """``pto.tmov ins(src) outs(dst)`` – move data between tile domains.""" kwargs = {} + if fp is not None: + kwargs["fp"] = unwrap_surface_value(fp) if mode is not None: kwargs["accToVecMode"] = _normalize_acc_to_vec_mode(mode, context="tmov(..., mode=...)") _pto.TMovOp(None, unwrap_surface_value(src), unwrap_surface_value(dst), **kwargs) @@ -2953,23 +2960,39 @@ def ttrans(src, tmp, dst): ) -def textract(src, dst, index_row, index_col): +def textract(src, dst, index_row, index_col, *, fp=None, mode=None): """``pto.textract ins(src, index_row, index_col) outs(dst)``.""" + kwargs = {} + if fp is not None: + kwargs["fp"] = unwrap_surface_value(fp) + if mode is not None: + kwargs["accToVecMode"] = _normalize_acc_to_vec_mode( + mode, context="textract(..., mode=...)" + ) _pto.TExtractOp( unwrap_surface_value(src), _coerce_index(index_row, context="textract(index_row)"), _coerce_index(index_col, context="textract(index_col)"), unwrap_surface_value(dst), + **kwargs, ) -def tinsert(src, dst, index_row, index_col): +def tinsert(src, dst, index_row, index_col, *, fp=None, mode=None): """``pto.tinsert ins(src, index_row, index_col) outs(dst)``.""" + kwargs = {} + if fp is not None: + kwargs["fp"] = unwrap_surface_value(fp) + if mode is not None: + kwargs["accToVecMode"] = _normalize_acc_to_vec_mode( + mode, context="tinsert(..., mode=...)" + ) _pto.TInsertOp( unwrap_surface_value(src), _coerce_index(index_row, context="tinsert(index_row)"), _coerce_index(index_col, context="tinsert(index_col)"), unwrap_surface_value(dst), + **kwargs, ) diff --git a/ptodsl/ptodsl/_tile_namespace.py b/ptodsl/ptodsl/_tile_namespace.py index 72995ff371..d12655c8c4 100644 --- a/ptodsl/ptodsl/_tile_namespace.py +++ b/ptodsl/ptodsl/_tile_namespace.py @@ -81,9 +81,9 @@ def load(src, tile, *, offsets=None, sizes=None): return _ops.tload(part, tile) @staticmethod - def store(tile, dst, *, offsets=None, sizes=None): + def store(tile, dst, *, offsets=None, sizes=None, fp=None): if offsets is None and sizes is None and _ops._is_partition_tensor_view(dst): - return _ops.tstore(tile, dst) + return _ops.tstore(tile, dst, fp=fp) part = _ops._tile_transfer_partition( dst, tile, @@ -91,7 +91,7 @@ def store(tile, dst, *, offsets=None, sizes=None): sizes=sizes, context="tile.store(...)", ) - return _ops.tstore(tile, part) + return _ops.tstore(tile, part, fp=fp) add = staticmethod(_ops.tadd) addrelu = staticmethod(_ops.taddrelu) diff --git a/ptodsl/tests/test_tilelib_catalog.py b/ptodsl/tests/test_tilelib_catalog.py index 24c2c1e9c1..10f68a6780 100644 --- a/ptodsl/tests/test_tilelib_catalog.py +++ b/ptodsl/tests/test_tilelib_catalog.py @@ -49,13 +49,6 @@ "pto.tcolsum": ("template_tcolsum", "pto.vadd", ("src", "dst"), "f32"), "pto.texpands": ("template_texpands", "pto.vdup", ("scalar", "dst"), "f32"), "pto.textract": ("template_textract_vec2vec_nd", "pto.vlds", ("src", "index_row", "index_col", "dst"), "f32"), - "pto.textract_fp": ( - "template_textract_fp_f32_f16", - "pto.mte_l0c_l1", - ("src", "fp", "index_row", "index_col", "dst"), - "f32", - "template_textract_fp_f32_f16", - ), "pto.tlrelu": ("template_tlrelu", "pto.vlrelu", ("src", "slope", "dst"), "f32"), "pto.tlog": ("template_tlog", "pto.vln", ("src", "dst"), "f32"), "pto.tdiv": ("template_tdiv", "pto.vdiv", ("src0", "src1", "dst"), "f32"), @@ -177,13 +170,6 @@ "f32", "template_tstore_nd", ), - "pto.tstore_fp": ( - "template_tstore_fp_acc_to_gm", - "pto.mte_l0c_gm", - ("src", "fp", "dst"), - "f32", - "template_tstore_fp_acc_to_gm", - ), "pto.tadds": ("template_tadds", "pto.vadds", ("src", "scalar", "dst"), "f32"), "pto.tmaxs": ("template_tmaxs", "pto.vmaxs", ("src", "scalar", "dst"), "f32"), "pto.tmins": ("template_tmins", "pto.vmins", ("src", "scalar", "dst"), "f32"), @@ -266,13 +252,12 @@ "pto.tilelang.instance", ) OPS_WITHOUT_TILE_LOAD = {"pto.texpands"} -OPS_WITHOUT_TILE_LOAD = OPS_WITHOUT_TILE_LOAD | {"pto.trandom", "pto.tsort32", "pto.tload", "pto.tstore", "pto.tstore_fp", "pto.textract_fp"} +OPS_WITHOUT_TILE_LOAD = OPS_WITHOUT_TILE_LOAD | {"pto.trandom", "pto.tsort32", "pto.tload", "pto.tstore"} OPS_WITHOUT_TILE_LOAD = OPS_WITHOUT_TILE_LOAD | CUBE_OPS OPS_WITHOUT_VECTOR_STORE = {"pto.tcmp", "pto.tcmps", "pto.tsort32"} -OPS_WITHOUT_VECTOR_STORE = OPS_WITHOUT_VECTOR_STORE | {"pto.tload", "pto.tstore", "pto.tstore_fp", "pto.textract_fp"} +OPS_WITHOUT_VECTOR_STORE = OPS_WITHOUT_VECTOR_STORE | {"pto.tload", "pto.tstore"} OPS_WITHOUT_VECTOR_STORE = OPS_WITHOUT_VECTOR_STORE | CUBE_OPS OPS_WITHOUT_LOOP = {"pto.tmrgsort"} -OPS_WITHOUT_LOOP = OPS_WITHOUT_LOOP | {"pto.tstore_fp", "pto.textract_fp"} OPS_WITHOUT_LOOP = OPS_WITHOUT_LOOP | CUBE_OPS OPS_ALLOWING_CASTPTR = {"pto.tsel", "pto.tsels"} SCALAR_OPERANDS = { @@ -293,8 +278,6 @@ ("pto.tshrs", "scalar"): "i16", ("pto.textract", "index_row"): "i32", ("pto.textract", "index_col"): "i32", - ("pto.textract_fp", "index_row"): "i32", - ("pto.textract_fp", "index_col"): "i32", ("pto.tinsert", "index_row"): "i32", ("pto.tinsert", "index_col"): "i32", ("pto.trandom", "key0"): "i32", @@ -309,10 +292,6 @@ ("pto.tcmps", "dst"): "ui8", ("pto.trandom", "dst"): "ui32", ("pto.tsort32", "idx"): "ui32", - ("pto.textract_fp", "fp"): "f32", - ("pto.textract_fp", "dst"): "f16", - ("pto.tstore_fp", "fp"): "f16", - ("pto.tstore_fp", "dst"): "f16", ("pto.trowargmax", "dst"): "i32", ("pto.trowargmin", "dst"): "i32", ("pto.tdequant", "scale"): "f32", @@ -347,17 +326,14 @@ VIEW_OPERANDS = { ("pto.tload", "src"), ("pto.tstore", "dst"), - ("pto.tstore_fp", "dst"), } VIEW_SHAPES = { ("pto.tload", "src"): (1, 1, 1, 8, 64), ("pto.tstore", "dst"): (1, 1, 1, 8, 64), - ("pto.tstore_fp", "dst"): (1, 1, 1, 8, 64), } VIEW_STRIDES = { ("pto.tload", "src"): (512, 512, 512, 64, 1), ("pto.tstore", "dst"): (512, 512, 512, 64, 1), - ("pto.tstore_fp", "dst"): (512, 512, 512, 64, 1), } for _op in CUBE_OPS: SPECIAL_MEMORY_SPACES[(_op, "lhs")] = "left" @@ -368,11 +344,6 @@ SPECIAL_MEMORY_SPACES[(_op, "bias")] = "bias" SPECIAL_MEMORY_SPACES[(_op, "lhs_scale")] = "scaling" SPECIAL_MEMORY_SPACES[(_op, "rhs_scale")] = "scaling" -SPECIAL_MEMORY_SPACES[("pto.textract_fp", "src")] = "acc" -SPECIAL_MEMORY_SPACES[("pto.textract_fp", "fp")] = "scaling" -SPECIAL_MEMORY_SPACES[("pto.textract_fp", "dst")] = "mat" -SPECIAL_MEMORY_SPACES[("pto.tstore_fp", "src")] = "acc" -SPECIAL_MEMORY_SPACES[("pto.tstore_fp", "fp")] = "scaling" for _op in ("pto.tgemv", "pto.tgemv.acc", "pto.tgemv.bias", "pto.tgemv.mx", "pto.tgemv.mx.acc", "pto.tgemv.mx.bias"): SPECIAL_VALID_SHAPES[(_op, "lhs")] = (1, 64) @@ -1328,7 +1299,7 @@ def test_tstore_nd_accepts_low_precision_tiles(self): self.assertEqual(selected.name, "template_tstore_nd") self.assertIn("pto.mte_ub_gm", selected.specialize(**specs).mlir_text()) - def test_textract_fp_versions_render(self): + def test_textract_fp_forms_use_unified_op(self): signatures = { ("f32", "f32", "i32", "i32", "si8"): "template_textract_fp_f32_si8", ("f32", "f32", "i32", "i32", "ui8"): "template_textract_fp_f32_ui8", @@ -1366,10 +1337,35 @@ def test_textract_fp_versions_render(self): memory_space="mat", ), } - selected = select("pto.textract_fp", "a5", specs) + selected = select("pto.textract", "a5", specs) self.assertEqual(selected.name, expected_name) self.assertIn("pto.mte_l0c_l1", selected.specialize(**specs).mlir_text()) + def test_tstore_fp_form_uses_unified_op(self): + specs = { + "src": TileSpec( + shape=(16, 32), + dtype=ScalarType("f32"), + memory_space="acc", + valid_shape=(16, 32), + b_layout="col_major", + s_layout="row_major", + ), + "dst": ViewSpec( + shape=(1, 1, 1, 16, 32), + dtype=ScalarType("f16"), + strides=(512, 512, 512, 32, 1), + ), + "fp": TileSpec( + shape=(1, 32), + dtype=ScalarType("f16"), + memory_space="scaling", + ), + } + selected = select("pto.tstore", "a5", specs) + self.assertEqual(selected.name, "template_tstore_fp_acc_to_gm") + self.assertIn("pto.mte_l0c_gm", selected.specialize(**specs).mlir_text()) + def test_tinsert_acc_to_mat_basic_renders(self): for dst_dtype in ("f16", "bf16"): with self.subTest(dst_dtype=dst_dtype): diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index 41dcaa0819..bc6cca0644 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -704,6 +704,29 @@ def test_tile_extract_dispatches_row_and_col_indices(self): (src, "idx:textract(index_row):7", "idx:textract(index_col):11", dst), ) self.assertEqual(coerce_index.call_count, 2) + + def test_tile_fp_forms_dispatch_through_unified_ops(self): + src = object() + dst = object() + fp = object() + + with patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "_is_partition_tensor_view", return_value=True), \ + patch.object(_ops, "_coerce_index", side_effect=lambda value, *, context: value), \ + patch.object(_ops._pto, "TStoreOp") as tstore_op, \ + patch.object(_ops._pto, "TMovOp") as tmov_op, \ + patch.object(_ops._pto, "TExtractOp") as textract_op, \ + patch.object(_ops._pto, "TInsertOp") as tinsert_op: + pto.tile.store(src, dst, fp=fp) + pto.tile.mov(src, dst, fp=fp) + pto.tile.extract(src, dst, 3, 5, fp=fp) + pto.tile.insert(src, dst, 3, 5, fp=fp) + + tstore_op.assert_called_once_with(None, src, dst, fp=fp) + tmov_op.assert_called_once_with(None, src, dst, fp=fp) + textract_op.assert_called_once_with(src, 3, 5, dst, fp=fp) + tinsert_op.assert_called_once_with(src, 3, 5, dst, fp=fp) + def test_sync_event_id_rejects_out_of_range_static_values(self): cases = [ (_ops.set_flag, ("MTE2", "V"), {"event_id": 8}, "set_flag(..., event_id=...)"), diff --git a/test/lit/pto/cube_tile_ops_positive.pto b/test/lit/pto/cube_tile_ops_positive.pto index b0ed6c9e4f..707ad68b36 100644 --- a/test/lit/pto/cube_tile_ops_positive.pto +++ b/test/lit/pto/cube_tile_ops_positive.pto @@ -99,7 +99,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind %tv = pto.make_tensor_view %dst_gm, shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view %sv = pto.partition_view %tv, offsets = [%c0, %c0], sizes = [%c1, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf16> - pto.tstore_fp ins(%acc_tile, %fp_tile : !pto.tile_buf, !pto.tile_buf) + pto.tstore ins(%acc_tile : !pto.tile_buf fp %fp_tile : !pto.tile_buf) outs(%sv : !pto.partition_tensor_view<1x32xf16>) pto.barrier #pto.pipe return diff --git a/test/lit/pto/extract_insert_fp_tile_native.pto b/test/lit/pto/extract_insert_fp_tile_native.pto index aa4e3f3f45..ac01ff22a2 100644 --- a/test/lit/pto/extract_insert_fp_tile_native.pto +++ b/test/lit/pto/extract_insert_fp_tile_native.pto @@ -18,7 +18,7 @@ module { %fp: !pto.tile_buf, %row: index, %col: index, %dst: !pto.tile_buf) { - pto.textract_fp ins(%src, %fp, %row, %col : !pto.tile_buf, !pto.tile_buf, index, index) + pto.textract ins(%src, %row, %col : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } @@ -28,16 +28,18 @@ module { %fp: !pto.tile_buf, %row: index, %col: index, %dst: !pto.tile_buf) { - pto.tinsert_fp ins(%src, %fp, %row, %col : !pto.tile_buf, !pto.tile_buf, index, index) + pto.tinsert ins(%src, %row, %col : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } } // NATIVE-LABEL: func.func private @textract_fp_arg( -// NATIVE: pto.textract_fp ins(%arg0, %arg1, %arg2, %arg3 +// NATIVE: pto.textract ins(%arg0, %arg2, %arg3 +// NATIVE-SAME: fp %arg1 // NATIVE-LABEL: func.func private @tinsert_fp_arg( -// NATIVE: pto.tinsert_fp ins(%arg0, %arg1, %arg2, %arg3 +// NATIVE: pto.tinsert ins(%arg0, %arg2, %arg3 +// NATIVE-SAME: fp %arg1 // NATIVE-NOT: memref< // EMITC-LABEL: textract_fp_arg( diff --git a/test/lit/pto/store_fp_tile_native.pto b/test/lit/pto/store_fp_tile_native.pto index 02397f0a2b..cf168b44ae 100644 --- a/test/lit/pto/store_fp_tile_native.pto +++ b/test/lit/pto/store_fp_tile_native.pto @@ -11,7 +11,7 @@ module { func.func private @tstore_fp_arg(%src: !pto.tile_buf, %fp: !pto.tile_buf, %dst: !pto.partition_tensor_view<1x32xf16>) { - pto.tstore_fp ins(%src, %fp : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x32xf16>) + pto.tstore ins(%src : !pto.tile_buf fp %fp : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x32xf16>) return } } @@ -20,7 +20,7 @@ module { // NATIVE-SAME: !pto.tile_buf -// NATIVE: pto.tstore_fp +// NATIVE: pto.tstore // EMITC-LABEL: tstore_fp_arg( // EMITC: TSTORE_FP( diff --git a/test/lit/pto/textract_acc_to_vec_a5_emitc.pto b/test/lit/pto/textract_acc_to_vec_a5_emitc.pto index fb427695da..553d5da5c1 100644 --- a/test/lit/pto/textract_acc_to_vec_a5_emitc.pto +++ b/test/lit/pto/textract_acc_to_vec_a5_emitc.pto @@ -25,7 +25,7 @@ module attributes {"pto.device-spec" = "Ascend950"} { {accToVecMode = #pto.acc_to_vec_mode, reluPreMode = #pto} - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst2 : !pto.tile_buf) {accToVecMode = #pto.acc_to_vec_mode} return diff --git a/test/lit/pto/textract_forms_emitc.pto b/test/lit/pto/textract_forms_emitc.pto index 8f0aeae864..74c6d6ebea 100644 --- a/test/lit/pto/textract_forms_emitc.pto +++ b/test/lit/pto/textract_forms_emitc.pto @@ -18,7 +18,7 @@ module attributes {"pto.device-spec" = "Ascend950"} { outs(%dst_quant : !pto.tile_buf) {reluPreMode = #pto} - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst_quant : !pto.tile_buf) {reluPreMode = #pto} return diff --git a/test/lit/pto/textract_fp_a3_lowering.pto b/test/lit/pto/textract_fp_a3_lowering.pto index b899ff07b7..d0d66c06f5 100644 --- a/test/lit/pto/textract_fp_a3_lowering.pto +++ b/test/lit/pto/textract_fp_a3_lowering.pto @@ -6,7 +6,7 @@ module { %src = pto.alloc_tile : !pto.tile_buf %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } diff --git a/test/lit/pto/textract_tinsert_low_precision_a5_emitc.pto b/test/lit/pto/textract_tinsert_low_precision_a5_emitc.pto index 113714874b..ce03d73406 100644 --- a/test/lit/pto/textract_tinsert_low_precision_a5_emitc.pto +++ b/test/lit/pto/textract_tinsert_low_precision_a5_emitc.pto @@ -18,9 +18,9 @@ module { %fp = pto.alloc_tile : !pto.tile_buf %mat_dst0 = pto.alloc_tile : !pto.tile_buf %mat_dst1 = pto.alloc_tile : !pto.tile_buf - pto.textract_fp ins(%acc_src, %fp, %c0, %c0 : !pto.tile_buf, !pto.tile_buf, index, index) + pto.textract ins(%acc_src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%mat_dst0 : !pto.tile_buf) - pto.tinsert_fp ins(%acc_src, %fp, %c0, %c0 : !pto.tile_buf, !pto.tile_buf, index, index) + pto.tinsert ins(%acc_src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%mat_dst1 : !pto.tile_buf) return } diff --git a/test/lit/pto/tinsert_a5_modes_emitc.pto b/test/lit/pto/tinsert_a5_modes_emitc.pto index 911dbfdb81..95e1bb8151 100644 --- a/test/lit/pto/tinsert_a5_modes_emitc.pto +++ b/test/lit/pto/tinsert_a5_modes_emitc.pto @@ -27,7 +27,7 @@ module attributes {"pto.device-spec" = "Ascend950"} { {accToVecMode = #pto.acc_to_vec_mode, reluPreMode = #pto} - pto.tinsert_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, !pto.tile_buf, index, index) + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst2 : !pto.tile_buf) {accToVecMode = #pto.acc_to_vec_mode} diff --git a/test/lit/pto/tinsert_forms_emitc.pto b/test/lit/pto/tinsert_forms_emitc.pto index 911bc89af4..bd00c6262b 100644 --- a/test/lit/pto/tinsert_forms_emitc.pto +++ b/test/lit/pto/tinsert_forms_emitc.pto @@ -19,7 +19,7 @@ module attributes {"pto.device-spec" = "Ascend950"} { outs(%dst_quant : !pto.tile_buf) {reluPreMode = #pto} - pto.tinsert_fp ins(%src_quant, %fp, %c0, %c0 : !pto.tile_buf, !pto.tile_buf, index, index) + pto.tinsert ins(%src_quant, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst_quant : !pto.tile_buf) {reluPreMode = #pto} return diff --git a/test/lit/pto/tinsert_fp_a3_lowering.pto b/test/lit/pto/tinsert_fp_a3_lowering.pto index 4f920caac3..f9b704782d 100644 --- a/test/lit/pto/tinsert_fp_a3_lowering.pto +++ b/test/lit/pto/tinsert_fp_a3_lowering.pto @@ -6,7 +6,7 @@ module { %src = pto.alloc_tile : !pto.tile_buf %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tinsert_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, !pto.tile_buf, index, index) + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } diff --git a/test/lit/pto/tinsert_fp_scaling_tile_role_emitc.pto b/test/lit/pto/tinsert_fp_scaling_tile_role_emitc.pto index fc589d0854..e68f6d3b3a 100644 --- a/test/lit/pto/tinsert_fp_scaling_tile_role_emitc.pto +++ b/test/lit/pto/tinsert_fp_scaling_tile_role_emitc.pto @@ -6,7 +6,7 @@ module { %src = pto.alloc_tile : !pto.tile_buf %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tinsert_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, !pto.tile_buf, index, index) + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } diff --git a/test/lit/pto/tmov_fp_tile_native.pto b/test/lit/pto/tmov_fp_tile_native.pto index 8cdcb77cca..16aca1576b 100644 --- a/test/lit/pto/tmov_fp_tile_native.pto +++ b/test/lit/pto/tmov_fp_tile_native.pto @@ -17,14 +17,15 @@ module { %src: !pto.tile_buf, %fp: !pto.tile_buf, %dst: !pto.tile_buf) { - pto.tmov.fp ins(%src, %fp : !pto.tile_buf, !pto.tile_buf) + pto.tmov ins(%src : !pto.tile_buf, %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } } // NATIVE-LABEL: func.func private @tmov_fp_arg( -// NATIVE: pto.tmov.fp ins(%arg0, %arg1 +// NATIVE: pto.tmov ins(%arg0 : +// NATIVE-SAME: %arg1 : // NATIVE-NOT: memref< // EMITC-LABEL: tmov_fp_arg( diff --git a/test/lit/pto/tstore_fp_insert_sync_effects.pto b/test/lit/pto/tstore_fp_insert_sync_effects.pto index 5803b9bd16..cde02edc73 100644 --- a/test/lit/pto/tstore_fp_insert_sync_effects.pto +++ b/test/lit/pto/tstore_fp_insert_sync_effects.pto @@ -13,7 +13,7 @@ // RUN: --mlir-print-ir-after=pto-insert-sync %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=WRITE -// `tstore_fp` writes its GM destination. A following tload reads the same GM +// The `tstore` fp form writes its GM destination. A following tload reads the same GM // partition, so insert-sync must emit a FIX -> MTE2 dependency after TSTORE_FP. module attributes {pto.target_arch = "a2a3"} { @@ -25,7 +25,7 @@ module attributes {pto.target_arch = "a2a3"} { %c0_i64 = arith.constant 0 : i64 %tmp = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf - pto.tstore_fp ins(%src, %fp : !pto.tile_buf, !pto.tile_buf) + pto.tstore ins(%src : !pto.tile_buf fp %fp : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x32xf16>) pto.tload ins(%dst : !pto.partition_tensor_view<1x32xf16>) outs(%tmp : !pto.tile_buf) @@ -44,7 +44,7 @@ module attributes {pto.target_arch = "a2a3"} { pto.tmatmul ins(%lhs, %rhs : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) - pto.tstore_fp ins(%acc, %fp : !pto.tile_buf, !pto.tile_buf) + pto.tstore ins(%acc : !pto.tile_buf fp %fp : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x32xf16>) return } @@ -54,10 +54,10 @@ module attributes {pto.target_arch = "a2a3"} { // READ: pto.tmatmul // READ: pto.set_flag[, , // READ: pto.wait_flag[, , -// READ: pto.tstore_fp +// READ: pto.tstore // WRITE-LABEL: func.func @tstore_fp_writes_dst -// WRITE: pto.tstore_fp +// WRITE: pto.tstore // WRITE: pto.set_flag[, , // WRITE: pto.wait_flag[, , // WRITE: pto.tload diff --git a/test/lit/pto/tstore_fp_invalid_dtype.pto b/test/lit/pto/tstore_fp_invalid_dtype.pto index 07322b3bac..27f34beb48 100644 --- a/test/lit/pto/tstore_fp_invalid_dtype.pto +++ b/test/lit/pto/tstore_fp_invalid_dtype.pto @@ -19,11 +19,11 @@ module attributes {pto.target_arch = "a5"} { %fp_tile = pto.alloc_tile : !pto.tile_buf %tv = pto.make_tensor_view %dst_gm, shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view %sv = pto.partition_view %tv, offsets = [%c0, %c0], sizes = [%c1, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf16> - pto.tstore_fp ins(%acc_tile, %fp_tile : !pto.tile_buf, !pto.tile_buf) + pto.tstore ins(%acc_tile : !pto.tile_buf fp %fp_tile : !pto.tile_buf) outs(%sv : !pto.partition_tensor_view<1x32xf16>) pto.barrier #pto.pipe return } } -// CHECK: expects src to have element type f32, i32 +// CHECK: expects A5 acc tstore src element type to be i32 or f32 diff --git a/test/lit/vpto/cube/expand_tile_op_tilelang_textract_fp_acc2mat.pto b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_fp_acc2mat.pto index 20ee4c110b..8c74955445 100644 --- a/test/lit/vpto/cube/expand_tile_op_tilelang_textract_fp_acc2mat.pto +++ b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_fp_acc2mat.pto @@ -11,7 +11,7 @@ // CHECK-LABEL: func.func @TEXTRACT_FP_A2M -// CHECK-NOT: pto.textract_fp +// CHECK-NOT: pto.textract // CHECK: pto.set_fpc @@ -45,13 +45,11 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, + blayout=col_major, slayout=row_major, fractal=1024, pad=0>, index, index fp %fp : !pto.tile_buf, index, index) + blayout=col_major, slayout=row_major, fractal=1024, pad=0>) outs(%dst : !pto.tile_buf} { @@ -30,8 +30,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, - !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } @@ -41,8 +40,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, - !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } @@ -52,8 +50,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, - !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } @@ -63,8 +60,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, - !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } @@ -74,8 +70,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, - !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } @@ -85,8 +80,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, - !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } @@ -96,8 +90,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, - !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } @@ -107,8 +100,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, - !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } @@ -118,8 +110,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, - !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } @@ -129,8 +120,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, - !pto.tile_buf, index, index) + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } diff --git a/test/lit/vpto/cube/textract_fp_verify_invalid_dst_loc.pto b/test/lit/vpto/cube/textract_fp_verify_invalid_dst_loc.pto index 787f69be8b..74b5b39c41 100644 --- a/test/lit/vpto/cube/textract_fp_verify_invalid_dst_loc.pto +++ b/test/lit/vpto/cube/textract_fp_verify_invalid_dst_loc.pto @@ -9,7 +9,7 @@ -// CHECK: expects dst to use loc=mat +// CHECK: expects A5 textract to use a supported src/dst loc pair @@ -39,13 +39,11 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, + blayout=col_major, slayout=row_major, fractal=1024, pad=0>, index, index fp %fp : !pto.tile_buf, index, index) + blayout=col_major, slayout=row_major, fractal=1024, pad=0>) outs(%dst : !pto.tile_buf, + blayout=col_major, slayout=row_major, fractal=1024, pad=0>, index, index fp %fp : !pto.tile_buf, index, index) + blayout=col_major, slayout=row_major, fractal=1024, pad=0>) outs(%dst : !pto.tile_buf, + blayout=col_major, slayout=row_major, fractal=1024, pad=0>, index, index fp %fp : !pto.tile_buf, index, index) + blayout=col_major, slayout=row_major, fractal=512, pad=0>) outs(%dst : !pto.tile_buf, + blayout=row_major, slayout=col_major, fractal=1024, pad=0>, index, index fp %fp : !pto.tile_buf, index, index) + blayout=col_major, slayout=row_major, fractal=1024, pad=0>) outs(%dst : !pto.tile_buf, + blayout=col_major, slayout=row_major, fractal=512, pad=0>, index, index fp %fp : !pto.tile_buf, index, index) + blayout=col_major, slayout=row_major, fractal=1024, pad=0>) outs(%dst : !pto.tile_bufMAT using fp (SCALING) tile. pto.TGemvOp(None, a_tile, b_tile, acc_tile) - pto.TMovFPOp(acc_tile, fp_scaling, out_mat) + pto.TMovOp(None, acc_tile, out_mat, fp=fp_scaling) pto.TStoreOp(None, acc_tile, sv_out) diff --git a/test/samples/Storefp/storefp.py b/test/samples/Storefp/storefp.py index a37f689366..a8d9887123 100644 --- a/test/samples/Storefp/storefp.py +++ b/test/samples/Storefp/storefp.py @@ -64,7 +64,7 @@ def build(): acc_tile = pto.AllocTileOp(acc_tile_ty).result fp_tile = pto.AllocTileOp(fp_tile_ty).result - pto.TStoreFPOp(acc_tile, fp_tile, sv) + pto.TStoreOp(None, acc_tile, sv, fp=fp_tile) func.ReturnOp([]) m.operation.verify() diff --git a/test/samples/Storefp/storefp_invalid.py b/test/samples/Storefp/storefp_invalid.py index aeed23011e..86b480d39e 100644 --- a/test/samples/Storefp/storefp_invalid.py +++ b/test/samples/Storefp/storefp_invalid.py @@ -56,7 +56,7 @@ def build(): dst = entry.arguments[0] src_tile = pto.AllocTileOp(src_tile_ty).result fp_tile = pto.AllocTileOp(fp_tile_ty).result - pto.TStoreFPOp(src_tile, fp_tile, dst) + pto.TStoreOp(None, src_tile, dst, fp=fp_tile) func.ReturnOp([]) ok = m.operation.verify() diff --git a/test/samples/TInsert/tinsert_fp.py b/test/samples/TInsert/tinsert_fp.py index 470dbf7beb..42f8506b71 100644 --- a/test/samples/TInsert/tinsert_fp.py +++ b/test/samples/TInsert/tinsert_fp.py @@ -64,7 +64,7 @@ def build(): src = pto.AllocTileOp(src_ty).result fp = pto.AllocTileOp(fp_ty).result dst = pto.AllocTileOp(dst_ty).result - pto.TInsertFPOp(src, fp, c0, c0, dst) + pto.TInsertOp(src, c0, c0, dst, fp=fp) func.ReturnOp([]) m.operation.verify() diff --git a/test/samples/runop.sh b/test/samples/runop.sh index 186ba63f2f..2db16abeda 100755 --- a/test/samples/runop.sh +++ b/test/samples/runop.sh @@ -1036,7 +1036,7 @@ PY if [[ "$base" == "extract_fp" ]]; then if ! grep -Fq "TEXTRACT_FP(" "$cpp"; then - echo -e "${A}(${base}.py)\tFAIL\tmissing TEXTRACT_FP() lowering for pto.textract_fp" + echo -e "${A}(${base}.py)\tFAIL\tmissing TEXTRACT_FP() lowering for pto.textract fp form" overall=1 continue fi @@ -1044,7 +1044,7 @@ PY if [[ "$base" == "tinsert_fp" ]]; then if ! grep -Fq "TINSERT_FP(" "$cpp"; then - echo -e "${A}(${base}.py)\tFAIL\tmissing TINSERT_FP() lowering for pto.tinsert_fp" + echo -e "${A}(${base}.py)\tFAIL\tmissing TINSERT_FP() lowering for pto.tinsert fp form" overall=1 continue fi diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/textract_fp/textract_fp.pto b/test/tilelang_st/npu/a5/src/st/smoke/testcase/textract_fp/textract_fp.pto index 7b21f47568..145799b05d 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/textract_fp/textract_fp.pto +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/textract_fp/textract_fp.pto @@ -132,9 +132,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind mat(f16) with pre_quant --- pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID0"] pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID0"] - pto.textract_fp ins(%acc_tile, %fp_tile, %c0_index, %c0_index : - !pto.tile_buf, - !pto.tile_buf, index, index) + pto.textract ins(%acc_tile, %c0_index, %c0_index : !pto.tile_buf, index, index fp %fp_tile : !pto.tile_buf) outs(%dst_mat : !pto.tile_buf) // --- Readback: textract_fp output x identity -> acc -> GM --- diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/textract_fp.pto b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/textract_fp.pto index 7b21f47568..145799b05d 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/textract_fp.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/textract_fp.pto @@ -132,9 +132,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind mat(f16) with pre_quant --- pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID0"] pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID0"] - pto.textract_fp ins(%acc_tile, %fp_tile, %c0_index, %c0_index : - !pto.tile_buf, - !pto.tile_buf, index, index) + pto.textract ins(%acc_tile, %c0_index, %c0_index : !pto.tile_buf, index, index fp %fp_tile : !pto.tile_buf) outs(%dst_mat : !pto.tile_buf) // --- Readback: textract_fp output x identity -> acc -> GM --- diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/tstore_acc2gm.pto b/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/tstore_acc2gm.pto index 6b6015f085..b81865161b 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/tstore_acc2gm.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/tstore_acc2gm.pto @@ -386,7 +386,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind GM f16 (qf322f16_pre_vec) %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view %sv_dst = pto.partition_view %tv_dst, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xf16> - pto.tstore_fp ins(%acc_c, %fp_tile : !pto.tile_buf, !pto.tile_buf) + pto.tstore ins(%acc_c : !pto.tile_buf fp %fp_tile : !pto.tile_buf) outs(%sv_dst : !pto.partition_tensor_view<16x32xf16>) pto.barrier #pto.pipe @@ -430,7 +430,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind GM bf16 (qf322bf16_pre_vec) %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view %sv_dst = pto.partition_view %tv_dst, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xbf16> - pto.tstore_fp ins(%acc_c, %fp_tile : !pto.tile_buf, !pto.tile_buf) + pto.tstore ins(%acc_c : !pto.tile_buf fp %fp_tile : !pto.tile_buf) outs(%sv_dst : !pto.partition_tensor_view<16x32xbf16>) pto.barrier #pto.pipe diff --git a/tools/ptobc/generated/ptobc_opcodes_v0.h b/tools/ptobc/generated/ptobc_opcodes_v0.h index ec0a1df46f..8072466e26 100644 --- a/tools/ptobc/generated/ptobc_opcodes_v0.h +++ b/tools/ptobc/generated/ptobc_opcodes_v0.h @@ -94,8 +94,7 @@ inline constexpr OpInfo kOpTable[] = { {0x101E, "pto.tdivs", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x101F, "pto.texp", 0, 0x00, 0x00, 2, 0, 0, 0x00}, {0x1020, "pto.texpands", 0, 0x00, 0x00, 2, 0, 0, 0x00}, - {0x1021, "pto.textract", 0, 0x00, 0x00, 4, 0, 0, 0x00}, - {0x1022, "pto.textract_fp", 0, 0x00, 0x00, 5, 0, 0, 0x00}, + {0x1021, "pto.textract", 0, 0x00, 0x02, 0, 0, 0, 0x00}, {0x1023, "pto.tfillpad", 0, 0x00, 0x00, 2, 0, 0, 0x00}, {0x1026, "pto.tfmod", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x1027, "pto.tfmods", 0, 0x00, 0x00, 3, 0, 0, 0x00}, @@ -104,8 +103,7 @@ inline constexpr OpInfo kOpTable[] = { {0x102A, "pto.tgemv", 1, 0x00, 0x01, 0, 0, 0, 0x00}, {0x102B, "pto.tgetval", 0, 0x01, 0x00, 2, 1, 0, 0x00}, {0x102C, "pto.timg2col", 0, 0x00, 0x00, 2, 0, 0, 0x00}, - {0x102D, "pto.tinsert", 0, 0x00, 0x00, 4, 0, 0, 0x00}, - {0x102E, "pto.tinsert_fp", 0, 0x00, 0x00, 5, 0, 0, 0x00}, + {0x102D, "pto.tinsert", 0, 0x00, 0x02, 0, 0, 0, 0x00}, {0x102F, "pto.tload", 0, 0x00, 0x00, 2, 0, 0, 0x00}, {0x1030, "pto.tlog", 0, 0x00, 0x00, 2, 0, 0, 0x00}, {0x1031, "pto.tlrelu", 0, 0x00, 0x00, 3, 0, 0, 0x00}, @@ -115,8 +113,7 @@ inline constexpr OpInfo kOpTable[] = { {0x1035, "pto.tmaxs", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x1036, "pto.tmin", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x1037, "pto.tmins", 0, 0x00, 0x00, 3, 0, 0, 0x00}, - {0x1038, "pto.tmov", 0, 0x00, 0x00, 2, 0, 0, 0x00}, - {0x1039, "pto.tmov.fp", 0, 0x00, 0x00, 3, 0, 0, 0x00}, + {0x1038, "pto.tmov", 0, 0x00, 0x02, 0, 0, 0, 0x00}, {0x103A, "pto.tmrgsort", 0, 0x00, 0x02, 0, 0, 0, 0x00}, {0x103B, "pto.tmul", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x103C, "pto.tmuls", 0, 0x00, 0x00, 3, 0, 0, 0x00}, @@ -161,7 +158,6 @@ inline constexpr OpInfo kOpTable[] = { {0x1063, "pto.tsort32", 0, 0x00, 0x02, 0, 0, 0, 0x00}, {0x1064, "pto.tsqrt", 0, 0x00, 0x00, 2, 0, 0, 0x00}, {0x1065, "pto.tstore", 0, 0x00, 0x02, 0, 0, 0, 0x00}, - {0x1066, "pto.tstore_fp", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x1067, "pto.tsub", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x1068, "pto.tsubc", 0, 0x00, 0x00, 4, 0, 0, 0x00}, {0x1069, "pto.tsubs", 0, 0x00, 0x00, 3, 0, 0, 0x00}, @@ -302,7 +298,6 @@ inline std::optional lookupOpcodeByName(llvm::StringRef name) { .Case("pto.texp", 0x101F) .Case("pto.texpands", 0x1020) .Case("pto.textract", 0x1021) - .Case("pto.textract_fp", 0x1022) .Case("pto.tfillpad", 0x1023) .Case("pto.tfmod", 0x1026) .Case("pto.tfmods", 0x1027) @@ -312,7 +307,6 @@ inline std::optional lookupOpcodeByName(llvm::StringRef name) { .Case("pto.tgetval", 0x102B) .Case("pto.timg2col", 0x102C) .Case("pto.tinsert", 0x102D) - .Case("pto.tinsert_fp", 0x102E) .Case("pto.tload", 0x102F) .Case("pto.tlog", 0x1030) .Case("pto.tlrelu", 0x1031) @@ -323,7 +317,6 @@ inline std::optional lookupOpcodeByName(llvm::StringRef name) { .Case("pto.tmin", 0x1036) .Case("pto.tmins", 0x1037) .Case("pto.tmov", 0x1038) - .Case("pto.tmov.fp", 0x1039) .Case("pto.tmrgsort", 0x103A) .Case("pto.tmul", 0x103B) .Case("pto.tmuls", 0x103C) @@ -368,7 +361,6 @@ inline std::optional lookupOpcodeByName(llvm::StringRef name) { .Case("pto.tsort32", 0x1063) .Case("pto.tsqrt", 0x1064) .Case("pto.tstore", 0x1065) - .Case("pto.tstore_fp", 0x1066) .Case("pto.tsub", 0x1067) .Case("pto.tsubc", 0x1068) .Case("pto.tsubs", 0x1069) @@ -496,7 +488,6 @@ inline std::optional lookupOpcodeAndVariantByFullName(llvm::St .Case("pto.texp", OpcodeAndVariant{0x101F, 0, 0}) .Case("pto.texpands", OpcodeAndVariant{0x1020, 0, 0}) .Case("pto.textract", OpcodeAndVariant{0x1021, 0, 0}) - .Case("pto.textract_fp", OpcodeAndVariant{0x1022, 0, 0}) .Case("pto.tfillpad", OpcodeAndVariant{0x1023, 0, 0}) .Case("pto.tfmod", OpcodeAndVariant{0x1026, 0, 0}) .Case("pto.tfmods", OpcodeAndVariant{0x1027, 0, 0}) @@ -505,7 +496,6 @@ inline std::optional lookupOpcodeAndVariantByFullName(llvm::St .Case("pto.tgetval", OpcodeAndVariant{0x102B, 0, 0}) .Case("pto.timg2col", OpcodeAndVariant{0x102C, 0, 0}) .Case("pto.tinsert", OpcodeAndVariant{0x102D, 0, 0}) - .Case("pto.tinsert_fp", OpcodeAndVariant{0x102E, 0, 0}) .Case("pto.tload", OpcodeAndVariant{0x102F, 0, 0}) .Case("pto.tlog", OpcodeAndVariant{0x1030, 0, 0}) .Case("pto.tlrelu", OpcodeAndVariant{0x1031, 0, 0}) @@ -514,7 +504,6 @@ inline std::optional lookupOpcodeAndVariantByFullName(llvm::St .Case("pto.tmin", OpcodeAndVariant{0x1036, 0, 0}) .Case("pto.tmins", OpcodeAndVariant{0x1037, 0, 0}) .Case("pto.tmov", OpcodeAndVariant{0x1038, 0, 0}) - .Case("pto.tmov.fp", OpcodeAndVariant{0x1039, 0, 0}) .Case("pto.tmrgsort", OpcodeAndVariant{0x103A, 0, 0}) .Case("pto.tmul", OpcodeAndVariant{0x103B, 0, 0}) .Case("pto.tmuls", OpcodeAndVariant{0x103C, 0, 0}) @@ -559,7 +548,6 @@ inline std::optional lookupOpcodeAndVariantByFullName(llvm::St .Case("pto.tsort32", OpcodeAndVariant{0x1063, 0, 0}) .Case("pto.tsqrt", OpcodeAndVariant{0x1064, 0, 0}) .Case("pto.tstore", OpcodeAndVariant{0x1065, 0, 0}) - .Case("pto.tstore_fp", OpcodeAndVariant{0x1066, 0, 0}) .Case("pto.tsub", OpcodeAndVariant{0x1067, 0, 0}) .Case("pto.tsubc", OpcodeAndVariant{0x1068, 0, 0}) .Case("pto.tsubs", OpcodeAndVariant{0x1069, 0, 0}) diff --git a/tools/ptobc/testdata/tstore_fp_v0_roundtrip.pto b/tools/ptobc/testdata/tstore_fp_v0_roundtrip.pto index dab24f2324..2ed94f5ad9 100644 --- a/tools/ptobc/testdata/tstore_fp_v0_roundtrip.pto +++ b/tools/ptobc/testdata/tstore_fp_v0_roundtrip.pto @@ -15,7 +15,7 @@ module { %dst_part = pto.partition_view %dst_tv, offsets = [%c0, %c0], sizes = [%c32, %c32] : !pto.tensor_view<32x32xi8> -> !pto.partition_tensor_view<32x32xi8> %acc_tile = pto.alloc_tile : !pto.tile_buf %fp_tile = pto.alloc_tile : !pto.tile_buf - pto.tstore_fp ins(%acc_tile, %fp_tile : !pto.tile_buf, !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<32x32xi8>) + pto.tstore ins(%acc_tile : !pto.tile_buf fp %fp_tile : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<32x32xi8>) return } } diff --git a/tools/ptobc/tests/tstore_fp_v0_encode.sh b/tools/ptobc/tests/tstore_fp_v0_encode.sh index 7d30d739a2..1da36ad102 100755 --- a/tools/ptobc/tests/tstore_fp_v0_encode.sh +++ b/tools/ptobc/tests/tstore_fp_v0_encode.sh @@ -31,5 +31,6 @@ ROUNDTRIP="${OUT_DIR}/tstore_fp_v0_roundtrip.roundtrip.pto" "${PTOBC_BIN}" encode "${IN}" -o "${BC}" "${PTOBC_BIN}" decode "${BC}" -o "${ROUNDTRIP}" -grep -F "pto.tstore_fp ins(" "${ROUNDTRIP}" >/dev/null +grep -F "pto.tstore ins(" "${ROUNDTRIP}" >/dev/null +grep -F " fp " "${ROUNDTRIP}" >/dev/null grep -F "!pto.partition_tensor_view<32x32xi8>" "${ROUNDTRIP}" >/dev/null From cdedab5e7ae8473457ca8c110e546484ddd48c5e Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Tue, 4 Aug 2026 15:40:52 +0800 Subject: [PATCH 004/122] fix: preserve PTOBC v0 FP operand compatibility --- .github/scripts/update_pto_isa_pin.py | 106 +++++++++++++++++- .github/workflows/ci_sim.yml | 7 +- .github/workflows/update_pto_isa_pin.yml | 14 ++- ReleaseNotes.md | 13 +++ docs/no_npu_compile_only_guide_zh.md | 4 +- tools/ptobc/MAINTENANCE.md | 17 ++- tools/ptobc/generated/ptobc_opcodes_v0.h | 20 +++- tools/ptobc/src/mlir_encode.cpp | 99 ++++++++++++++++ tools/ptobc/src/ptobc_decode_print.cpp | 101 +++++++++++++++++ .../fp_operand_forms_v0_roundtrip.pto | 24 ++++ tools/ptobc/tests/CMakeLists.txt | 17 +++ .../ptobc/tests/fp_operand_forms_v0_encode.sh | 53 +++++++++ tools/ptobc/tests/tstore_fp_v0_encode.sh | 23 ++++ .../tests/v0_fp_schema_compatibility_check.py | 54 +++++++++ 14 files changed, 537 insertions(+), 15 deletions(-) create mode 100644 tools/ptobc/testdata/fp_operand_forms_v0_roundtrip.pto create mode 100755 tools/ptobc/tests/fp_operand_forms_v0_encode.sh create mode 100755 tools/ptobc/tests/v0_fp_schema_compatibility_check.py diff --git a/.github/scripts/update_pto_isa_pin.py b/.github/scripts/update_pto_isa_pin.py index 4bc661fb54..81ee204d6c 100755 --- a/.github/scripts/update_pto_isa_pin.py +++ b/.github/scripts/update_pto_isa_pin.py @@ -40,6 +40,16 @@ def parse_args() -> argparse.Namespace: default="docker/Dockerfile", help="Path to the Dockerfile that vendors pto-isa.", ) + parser.add_argument( + "--sim-workflow", + default=".github/workflows/ci_sim.yml", + help="Path to the VPTO simulator workflow file.", + ) + parser.add_argument( + "--compile-only-guide", + default="docs/no_npu_compile_only_guide_zh.md", + help="Path to the no-NPU compile-only guide.", + ) parser.add_argument( "--remote-validation-script", default="test/npu_validation/scripts/run_remote_npu_validation.sh", @@ -83,6 +93,21 @@ def replace_exactly_once( return new_text +def replace_exact_count( + text: str, + pattern: str, + replacement: str, + path: pathlib.Path, + expected_count: int, +) -> str: + new_text, count = re.subn(pattern, replacement, text, flags=re.MULTILINE) + if count != expected_count: + raise RuntimeError( + f"expected {expected_count} matches for pattern {pattern!r} in {path}, got {count}" + ) + return new_text + + def update_ci_workflow(path: pathlib.Path, commit: str) -> bool: original = read_text(path) updated = original @@ -125,6 +150,35 @@ def update_dockerfile(path: pathlib.Path, commit: str) -> bool: return False +def update_sim_workflow(path: pathlib.Path, commit: str) -> bool: + original = read_text(path) + updated = replace_exactly_once( + original, + r"^(\s*PTO_ISA_COMMIT:\s*)([0-9a-f]{40})$", + rf"\g<1>{commit}", + path, + ) + if updated != original: + write_text(path, updated) + return True + return False + + +def update_compile_only_guide(path: pathlib.Path, commit: str) -> bool: + original = read_text(path) + updated = replace_exact_count( + original, + r"^(export PTO_ISA_COMMIT=)([0-9a-f]{40})$", + rf"\g<1>{commit}", + path, + expected_count=2, + ) + if updated != original: + write_text(path, updated) + return True + return False + + def update_remote_validation_script(path: pathlib.Path, commit: str) -> bool: original = read_text(path) updated = replace_exactly_once( @@ -168,6 +222,29 @@ def extract_docker_commit(path: pathlib.Path) -> tuple[str, str]: return arg_match.group(1), comment_match.group(1) +def extract_sim_commit(path: pathlib.Path) -> str: + text = read_text(path) + match = re.search( + r"^\s*PTO_ISA_COMMIT:\s*([0-9a-f]{40})$", text, flags=re.MULTILINE + ) + if not match: + raise RuntimeError(f"failed to read pinned pto-isa commit from {path}") + return match.group(1) + + +def extract_compile_only_commits(path: pathlib.Path) -> tuple[str, str]: + matches = re.findall( + r"^export PTO_ISA_COMMIT=([0-9a-f]{40})$", + read_text(path), + flags=re.MULTILINE, + ) + if len(matches) != 2: + raise RuntimeError( + f"expected two pinned pto-isa commits in {path}, got {len(matches)}" + ) + return matches[0], matches[1] + + def extract_remote_validation_commit(path: pathlib.Path) -> str: text = read_text(path) match = re.search( @@ -183,11 +260,15 @@ def extract_remote_validation_commit(path: pathlib.Path) -> str: def verify( ci_path: pathlib.Path, docker_path: pathlib.Path, + sim_workflow_path: pathlib.Path, + compile_only_guide_path: pathlib.Path, remote_validation_path: pathlib.Path, commit: str, ) -> None: ci_default, ci_env = extract_ci_commit(ci_path) docker_arg, docker_comment = extract_docker_commit(docker_path) + sim_commit = extract_sim_commit(sim_workflow_path) + guide_setup, guide_run = extract_compile_only_commits(compile_only_guide_path) remote_validation_commit = extract_remote_validation_commit( remote_validation_path ) @@ -196,6 +277,9 @@ def verify( f"{ci_path}:runtime_default": ci_env, f"{docker_path}:arg": docker_arg, f"{docker_path}:comment": docker_comment, + f"{sim_workflow_path}:simulator": sim_commit, + f"{compile_only_guide_path}:setup": guide_setup, + f"{compile_only_guide_path}:run": guide_run, f"{remote_validation_path}:fallback": remote_validation_commit, } mismatches = {name: value for name, value in values.items() if value != commit} @@ -209,17 +293,35 @@ def main() -> int: commit = args.commit or resolve_head_commit(args.repo_url) ci_path = pathlib.Path(args.ci_workflow) docker_path = pathlib.Path(args.dockerfile) + sim_workflow_path = pathlib.Path(args.sim_workflow) + compile_only_guide_path = pathlib.Path(args.compile_only_guide) remote_validation_path = pathlib.Path(args.remote_validation_script) if args.check: - verify(ci_path, docker_path, remote_validation_path, commit) + verify( + ci_path, + docker_path, + sim_workflow_path, + compile_only_guide_path, + remote_validation_path, + commit, + ) print(commit) return 0 update_ci_workflow(ci_path, commit) update_dockerfile(docker_path, commit) + update_sim_workflow(sim_workflow_path, commit) + update_compile_only_guide(compile_only_guide_path, commit) update_remote_validation_script(remote_validation_path, commit) - verify(ci_path, docker_path, remote_validation_path, commit) + verify( + ci_path, + docker_path, + sim_workflow_path, + compile_only_guide_path, + remote_validation_path, + commit, + ) print(commit) return 0 diff --git a/.github/workflows/ci_sim.yml b/.github/workflows/ci_sim.yml index 2bf1d56d42..65b59983b9 100644 --- a/.github/workflows/ci_sim.yml +++ b/.github/workflows/ci_sim.yml @@ -116,10 +116,9 @@ jobs: PYPTO_REF: ef6ce7cd8bd33b4c93b58dc34830a20b145736ac PYPTO_WORKSPACE: ${{ github.workspace }}/.work/pypto-ci PYPTO_RUN_WORKSPACE: ${{ github.workspace }}/.work/pypto-run-ci - # GitHub pin containing the two-argument Soft SYNCALL ABI from f24f7b7 - # and the follow-up CPU-Sim duplicate-stub fix. The GitCode mirror carries - # the ABI as d56d42d within the separate ce3262e3 remote-validation pin. - PTO_ISA_COMMIT: a8168c6ca8ae22b06e4ea440506f0a87e446e9a2 + # Keep this pin aligned with the compiler, Docker, documentation, and + # remote-validation defaults via update_pto_isa_pin.py. + PTO_ISA_COMMIT: 27386d906e8fdcbd93aec84197939bc0b2c6caea PTO_ISA_ROOT: ${{ github.workspace }}/.work/pto-isa-ci steps: - name: Checkout diff --git a/.github/workflows/update_pto_isa_pin.yml b/.github/workflows/update_pto_isa_pin.yml index 13ab8046b4..15057f4e4d 100644 --- a/.github/workflows/update_pto_isa_pin.yml +++ b/.github/workflows/update_pto_isa_pin.yml @@ -60,7 +60,12 @@ jobs: shell: bash run: | set -euo pipefail - if git diff --quiet -- .github/workflows/ci.yml docker/Dockerfile; then + if git diff --quiet -- \ + .github/workflows/ci.yml \ + .github/workflows/ci_sim.yml \ + docker/Dockerfile \ + docs/no_npu_compile_only_guide_zh.md \ + test/npu_validation/scripts/run_remote_npu_validation.sh; then echo "pto-isa pin already up to date." echo "created_commit=false" >> "${GITHUB_OUTPUT}" exit 0 @@ -68,7 +73,12 @@ jobs: git config user.name "github-actions[bot]" git config user.email "41898282+github-actions[bot]@users.noreply.github.com" - git add .github/workflows/ci.yml docker/Dockerfile + git add \ + .github/workflows/ci.yml \ + .github/workflows/ci_sim.yml \ + docker/Dockerfile \ + docs/no_npu_compile_only_guide_zh.md \ + test/npu_validation/scripts/run_remote_npu_validation.sh git commit -m "chore(ci): bump pto-isa pin to ${{ steps.update.outputs.pto_isa_short }}" echo "created_commit=true" >> "${GITHUB_OUTPUT}" diff --git a/ReleaseNotes.md b/ReleaseNotes.md index 79882f63e0..74d7a92b0c 100644 --- a/ReleaseNotes.md +++ b/ReleaseNotes.md @@ -1,5 +1,18 @@ # PTOAS (PTO Assembler & Optimizer) +## 未发布 + +### PTO IR 接口调整 + +- `pto.tfillpad_inplace` 和 `pto.tfillpad_expand` 已整合到 + `pto.tfillpad`；通过 `mode` 选择行为，默认值为 `normal`。 +- `pto.textract_fp`、`pto.tinsert_fp`、`pto.tmov.fp` 和 + `pto.tstore_fp` 已移除。请改用 `pto.textract`、`pto.tinsert`、 + `pto.tmov` 和 `pto.tstore` 的可选 `fp` operand。旧文本 PTO IR + 不会自动迁移，需要在升级时替换 op 名称并按新语法传入 `fp`。 +- PTO-BC v0 保留上述 FP op 的历史 opcode 作为 wire 兼容别名；已有 + `.ptobc` 文件可由新版 decoder 读取并转换为统一后的 PTO IR。 + ## 版本 - 版本号：v0.51 - 发布日期：2026-02-14 diff --git a/docs/no_npu_compile_only_guide_zh.md b/docs/no_npu_compile_only_guide_zh.md index c96b7a1b35..344911e3bf 100644 --- a/docs/no_npu_compile_only_guide_zh.md +++ b/docs/no_npu_compile_only_guide_zh.md @@ -168,7 +168,7 @@ cmake --build build --parallel ```bash export PAYLOAD_ROOT=/tmp/ptoas_payload export TARGET_SOC_VERSION=Ascend910 -export PTO_ISA_COMMIT=ce3262e3825a235f951917eeada30e52910b6a84 +export PTO_ISA_COMMIT=27386d906e8fdcbd93aec84197939bc0b2c6caea rm -rf "$PAYLOAD_ROOT" mkdir -p "$PAYLOAD_ROOT/test/samples" @@ -220,7 +220,7 @@ export STAGE=build export RUN_MODE=npu export SOC_VERSION="$TARGET_SOC_VERSION" export PTO_ISA_REPO=https://gitcode.com/cann/pto-isa.git -export PTO_ISA_COMMIT=ce3262e3825a235f951917eeada30e52910b6a84 +export PTO_ISA_COMMIT=27386d906e8fdcbd93aec84197939bc0b2c6caea # 参照 CI 的做法，按目标 SoC 排除非匹配的 A3/A5 变体。 A3_ONLY_CASES="partition5d,partition5d_dynamic,mrgsort,tmatmulk_autosync" diff --git a/tools/ptobc/MAINTENANCE.md b/tools/ptobc/MAINTENANCE.md index b6a51129cc..3038e7f9a6 100644 --- a/tools/ptobc/MAINTENANCE.md +++ b/tools/ptobc/MAINTENANCE.md @@ -2,18 +2,33 @@ This tool encodes/decodes PTO-BC v0. +## v0 compatibility contract + +PTO-BC v0 files are expected to remain readable across PTOAS builds. Never +change the payload schema of an assigned opcode while keeping version 0. When +an IR op gains optional operands, preserve the shipped opcode payload and use +either a new opcode, a legacy wire alias, or the generic v0 compatibility +encoding for forms that do not fit the old schema. + +`tools/ptobc/generated/ptobc_opcodes_v0.h` is the checked-in authoritative +schema table. The generator named by older header comments is not present in +this repository, so table changes are currently maintained and reviewed by +hand. + ## When you change the PTO dialect / IR If you change any of the following: - `include/PTO/IR/PTOOps.td` (add/remove ops, rename mnemonics) - operand counts / region structure / immediates semantics -…then you **must** update the PTO-BC v0 opcode/schema tables (regenerate `tools/ptobc/generated/ptobc_opcodes_v0.h`) and ensure tests pass. +…then you **must** update the PTO-BC v0 opcode/schema table without changing +existing wire payloads and ensure tests pass. ## Required gates Run (or rely on CI): - `ctest -R ptobc_stage9_e2e` - `ctest -R ptobc_to_ptoas_smoke` - `ctest -R ptobc_opcode_coverage_check` +- `ctest -R ptobc_v0_fp_schema_compatibility_check` ## Notes - `ptobc_opcode_coverage_check` is a heuristic based on `mnemonic = "..."` occurrences. diff --git a/tools/ptobc/generated/ptobc_opcodes_v0.h b/tools/ptobc/generated/ptobc_opcodes_v0.h index 8072466e26..ac5b41bbe4 100644 --- a/tools/ptobc/generated/ptobc_opcodes_v0.h +++ b/tools/ptobc/generated/ptobc_opcodes_v0.h @@ -6,7 +6,11 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// Generated by docs/bytecode/tools/gen_v0_tables.py +// PTO-BC v0 schema table. +// +// This checked-in table is authoritative. The generator referenced by older +// revisions is not shipped in this repository, so changes must be reviewed +// against tools/ptobc/MAINTENANCE.md and the v0 compatibility tests. #pragma once #include @@ -94,7 +98,9 @@ inline constexpr OpInfo kOpTable[] = { {0x101E, "pto.tdivs", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x101F, "pto.texp", 0, 0x00, 0x00, 2, 0, 0, 0x00}, {0x1020, "pto.texpands", 0, 0x00, 0x00, 2, 0, 0, 0x00}, - {0x1021, "pto.textract", 0, 0x00, 0x02, 0, 0, 0, 0x00}, + {0x1021, "pto.textract", 0, 0x00, 0x00, 4, 0, 0, 0x00}, + // Legacy textract_fp wire opcode; decoded as the unified pto.textract op. + {0x1022, "pto.textract", 0, 0x00, 0x00, 5, 0, 0, 0x00}, {0x1023, "pto.tfillpad", 0, 0x00, 0x00, 2, 0, 0, 0x00}, {0x1026, "pto.tfmod", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x1027, "pto.tfmods", 0, 0x00, 0x00, 3, 0, 0, 0x00}, @@ -103,7 +109,9 @@ inline constexpr OpInfo kOpTable[] = { {0x102A, "pto.tgemv", 1, 0x00, 0x01, 0, 0, 0, 0x00}, {0x102B, "pto.tgetval", 0, 0x01, 0x00, 2, 1, 0, 0x00}, {0x102C, "pto.timg2col", 0, 0x00, 0x00, 2, 0, 0, 0x00}, - {0x102D, "pto.tinsert", 0, 0x00, 0x02, 0, 0, 0, 0x00}, + {0x102D, "pto.tinsert", 0, 0x00, 0x00, 4, 0, 0, 0x00}, + // Legacy tinsert_fp wire opcode; decoded as the unified pto.tinsert op. + {0x102E, "pto.tinsert", 0, 0x00, 0x00, 5, 0, 0, 0x00}, {0x102F, "pto.tload", 0, 0x00, 0x00, 2, 0, 0, 0x00}, {0x1030, "pto.tlog", 0, 0x00, 0x00, 2, 0, 0, 0x00}, {0x1031, "pto.tlrelu", 0, 0x00, 0x00, 3, 0, 0, 0x00}, @@ -113,7 +121,9 @@ inline constexpr OpInfo kOpTable[] = { {0x1035, "pto.tmaxs", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x1036, "pto.tmin", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x1037, "pto.tmins", 0, 0x00, 0x00, 3, 0, 0, 0x00}, - {0x1038, "pto.tmov", 0, 0x00, 0x02, 0, 0, 0, 0x00}, + {0x1038, "pto.tmov", 0, 0x00, 0x00, 2, 0, 0, 0x00}, + // Legacy tmov.fp wire opcode; decoded as the unified pto.tmov op. + {0x1039, "pto.tmov", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x103A, "pto.tmrgsort", 0, 0x00, 0x02, 0, 0, 0, 0x00}, {0x103B, "pto.tmul", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x103C, "pto.tmuls", 0, 0x00, 0x00, 3, 0, 0, 0x00}, @@ -158,6 +168,8 @@ inline constexpr OpInfo kOpTable[] = { {0x1063, "pto.tsort32", 0, 0x00, 0x02, 0, 0, 0, 0x00}, {0x1064, "pto.tsqrt", 0, 0x00, 0x00, 2, 0, 0, 0x00}, {0x1065, "pto.tstore", 0, 0x00, 0x02, 0, 0, 0, 0x00}, + // Legacy tstore_fp wire opcode; decoded as the unified pto.tstore op. + {0x1066, "pto.tstore", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x1067, "pto.tsub", 0, 0x00, 0x00, 3, 0, 0, 0x00}, {0x1068, "pto.tsubc", 0, 0x00, 0x00, 4, 0, 0, 0x00}, {0x1069, "pto.tsubs", 0, 0x00, 0x00, 3, 0, 0, 0x00}, diff --git a/tools/ptobc/src/mlir_encode.cpp b/tools/ptobc/src/mlir_encode.cpp index a7c2914a8a..232a6cf23f 100644 --- a/tools/ptobc/src/mlir_encode.cpp +++ b/tools/ptobc/src/mlir_encode.cpp @@ -35,6 +35,7 @@ #include #include #include +#include #include namespace ptobc { @@ -55,6 +56,14 @@ constexpr uint8_t kCmpPredicateSltEncoding = 2; constexpr uint8_t kCmpPredicateSleEncoding = 3; constexpr uint8_t kCmpPredicateSgtEncoding = 4; constexpr uint8_t kCmpPredicateSgeEncoding = 5; +constexpr uint16_t kTExtractOpcode = 0x1021; +constexpr uint16_t kTExtractFpWireOpcode = 0x1022; +constexpr uint16_t kTInsertOpcode = 0x102D; +constexpr uint16_t kTInsertFpWireOpcode = 0x102E; +constexpr uint16_t kTMovOpcode = 0x1038; +constexpr uint16_t kTMovFpWireOpcode = 0x1039; +constexpr uint16_t kTStoreOpcode = 0x1065; +constexpr uint16_t kTStoreFpWireOpcode = 0x1066; using NamedAttributeVector = llvm::SmallVector; @@ -74,6 +83,15 @@ static bool shouldEncodeViaGenericV0CompatibilityShim(mlir::Operation &op) { return static_cast(tci.getTmp()); if (auto trowexpandadd = llvm::dyn_cast(&op)) return static_cast(trowexpandadd.getTmp()); + // The compact v0 schemas for these ops predate their optional pre-quant + // operands. Keep the shipped fixed payloads unchanged and use the generic + // opcode for forms that cannot be represented by those schemas. + if (auto textract = llvm::dyn_cast(&op)) + return static_cast(textract.getPreQuantScalar()); + if (auto tinsert = llvm::dyn_cast(&op)) + return static_cast(tinsert.getPreQuantScalar()); + if (auto tmov = llvm::dyn_cast(&op)) + return static_cast(tmov.getPreQuantScalar()); if (llvm::isa( &op)) return true; @@ -91,6 +109,39 @@ static bool shouldEncodeViaGenericV0CompatibilityShim(mlir::Operation &op) { return false; } +static std::optional +getLegacyFpWireOpcode(mlir::Operation &op) { + if (auto textract = llvm::dyn_cast(&op)) + return textract.getFp() ? std::optional(kTExtractFpWireOpcode) + : std::nullopt; + if (auto tinsert = llvm::dyn_cast(&op)) + return tinsert.getFp() ? std::optional(kTInsertFpWireOpcode) + : std::nullopt; + if (auto tmov = llvm::dyn_cast(&op)) + return tmov.getFp() ? std::optional(kTMovFpWireOpcode) + : std::nullopt; + if (auto tstore = llvm::dyn_cast(&op)) + return tstore.getFp() ? std::optional(kTStoreFpWireOpcode) + : std::nullopt; + return std::nullopt; +} + +static bool omitsDerivedOperandSegmentsInV0(uint16_t opcode) { + switch (opcode) { + case kTExtractOpcode: + case kTExtractFpWireOpcode: + case kTInsertOpcode: + case kTInsertFpWireOpcode: + case kTMovOpcode: + case kTMovFpWireOpcode: + case kTStoreOpcode: + case kTStoreFpWireOpcode: + return true; + default: + return false; + } +} + static uint64_t internType(PTOBCFile& f, mlir::Type t) { std::string s = printType(t); f.strings.intern(s); @@ -553,6 +604,41 @@ void Encoder::encodeKnownOpOperands( writeULEB128(getValueId(tscatter.getDst()), out.bytes); return true; }; + auto emitLegacyFpOperands = [&]() { + auto emit = [&](llvm::ArrayRef operands) { + for (mlir::Value value : operands) + writeULEB128(getValueId(value), out.bytes); + }; + switch (variantInfo.opcode) { + case kTExtractFpWireOpcode: { + auto fpOp = llvm::cast(&op); + emit({fpOp.getSrc(), fpOp.getFp(), fpOp.getIndexRow(), fpOp.getIndexCol(), + fpOp.getDst()}); + return true; + } + case kTInsertFpWireOpcode: { + auto fpOp = llvm::cast(&op); + emit({fpOp.getSrc(), fpOp.getFp(), fpOp.getIndexRow(), fpOp.getIndexCol(), + fpOp.getDst()}); + return true; + } + case kTMovFpWireOpcode: { + auto fpOp = llvm::cast(&op); + emit({fpOp.getSrc(), fpOp.getFp(), fpOp.getDst()}); + return true; + } + case kTStoreFpWireOpcode: { + auto fpOp = llvm::cast(&op); + emit({fpOp.getSrc(), fpOp.getFp(), fpOp.getDst()}); + return true; + } + default: + return false; + } + }; + + if (emitLegacyFpOperands()) + return; switch (info.operand_mode) { case 0x00: @@ -604,6 +690,8 @@ void Encoder::encodeKnownOp(mlir::Operation &op, Buffer &out, out.appendU16LE(variantInfo.opcode); mlir::DictionaryAttr dict = op.getAttrDictionary(); dict = stripKnownImmediateAttrs(op.getContext(), dict, info); + if (omitsDerivedOperandSegmentsInV0(variantInfo.opcode)) + dict = stripAttrs(op.getContext(), dict, {"operandSegmentSizes"}); // The assembly printer omits default pipe ids on transfer/free ops. Keep v0 // byte roundtrips stable by using the same representation while encoding. dict = dropDefaultZeroPipeIdForV0Encoding(op, dict); @@ -687,6 +775,17 @@ void Encoder::encodeOp(mlir::Operation& op, Buffer& out) { return; } + if (auto opcode = getLegacyFpWireOpcode(op)) { + auto variantInfo = + ptobc::v0::OpcodeAndVariant{*opcode, /*hasVariant=*/0, /*variant=*/0}; + const auto *info = ptobc::v0::lookupByOpcode(*opcode); + if (!info) + throw std::runtime_error("missing legacy FP v0 opcode schema for op: " + + fullName.str()); + encodeKnownOp(op, out, *info, variantInfo); + return; + } + auto variantInfo = ptobc::v0::lookupOpcodeAndVariantByFullName(fullName); if (variantInfo) { const auto *info = ptobc::v0::lookupByOpcode(variantInfo->opcode); diff --git a/tools/ptobc/src/ptobc_decode_print.cpp b/tools/ptobc/src/ptobc_decode_print.cpp index cdb0280902..1d394de732 100644 --- a/tools/ptobc/src/ptobc_decode_print.cpp +++ b/tools/ptobc/src/ptobc_decode_print.cpp @@ -57,6 +57,14 @@ constexpr size_t kPTOBCMagicSize = 6; constexpr size_t kPTOBCHeaderSize = 14; constexpr size_t kPTOBCVersionOffset = 6; constexpr size_t kPTOBCPayloadLengthOffset = 10; +constexpr uint16_t kTExtractOpcode = 0x1021; +constexpr uint16_t kTExtractFpWireOpcode = 0x1022; +constexpr uint16_t kTInsertOpcode = 0x102D; +constexpr uint16_t kTInsertFpWireOpcode = 0x102E; +constexpr uint16_t kTMovOpcode = 0x1038; +constexpr uint16_t kTMovFpWireOpcode = 0x1039; +constexpr uint16_t kTStoreOpcode = 0x1065; +constexpr uint16_t kTStoreFpWireOpcode = 0x1066; struct Reader { const uint8_t* p; @@ -507,6 +515,92 @@ materializeOperands(BuildCtx &bc, llvm::ArrayRef operandIds) { return operands; } +static void normalizeLegacyFpOperandOrder( + uint16_t opcode, + llvm::SmallVectorImpl &operands) { + if ((opcode == kTExtractFpWireOpcode || + opcode == kTInsertFpWireOpcode) && + operands.size() == 5) { + // Legacy wire order: src, fp, row, col, dst. + llvm::SmallVector reordered{ + operands[0], operands[2], operands[3], operands[4], operands[1]}; + operands.assign(reordered.begin(), reordered.end()); + return; + } + if ((opcode == kTMovFpWireOpcode || opcode == kTStoreFpWireOpcode) && + operands.size() == 3) { + // Legacy wire order: src, fp, dst. + llvm::SmallVector reordered{operands[0], operands[2], + operands[1]}; + operands.assign(reordered.begin(), reordered.end()); + } +} + +static std::optional> +getUnifiedOperandSegments(uint16_t opcode, + llvm::ArrayRef operands) { + switch (opcode) { + case kTExtractOpcode: + return llvm::SmallVector{1, 1, 1, 1, 0, 0}; + case kTExtractFpWireOpcode: + return llvm::SmallVector{1, 1, 1, 1, 1, 0}; + case kTInsertOpcode: + return llvm::SmallVector{1, 1, 1, 1, 0, 0}; + case kTInsertFpWireOpcode: + return llvm::SmallVector{1, 1, 1, 1, 1, 0}; + case kTMovOpcode: + return llvm::SmallVector{1, 1, 0, 0}; + case kTMovFpWireOpcode: + return llvm::SmallVector{1, 1, 1, 0}; + case kTStoreOpcode: + if (operands.size() == 2) + return llvm::SmallVector{1, 1, 0, 0}; + if (operands.size() == 3) { + const bool thirdIsFp = + mlir::isa(operands.back().getType()); + return llvm::SmallVector{1, 1, thirdIsFp ? 1 : 0, + thirdIsFp ? 0 : 1}; + } + return std::nullopt; + case kTStoreFpWireOpcode: + return llvm::SmallVector{1, 1, 1, 0}; + default: + return std::nullopt; + } +} + +static void setUnifiedOperandSegmentProperty( + mlir::Operation *op, uint16_t opcode, llvm::ArrayRef segments) { + switch (opcode) { + case kTExtractOpcode: + case kTExtractFpWireOpcode: + llvm::cast(op) + .getProperties() + .setOperandSegmentSizes(segments); + return; + case kTInsertOpcode: + case kTInsertFpWireOpcode: + llvm::cast(op) + .getProperties() + .setOperandSegmentSizes(segments); + return; + case kTMovOpcode: + case kTMovFpWireOpcode: + llvm::cast(op) + .getProperties() + .setOperandSegmentSizes(segments); + return; + case kTStoreOpcode: + case kTStoreFpWireOpcode: + llvm::cast(op) + .getProperties() + .setOperandSegmentSizes(segments); + return; + default: + return; + } +} + static mlir::Operation *buildGenericOpFromReader(BuildCtx &bc, Reader &r, mlir::Block &block, uint64_t opId, @@ -584,6 +678,7 @@ static mlir::Operation *buildKnownOpFromReader(BuildCtx &bc, Reader &r, KnownOpImmediates imms = readKnownOpImmediates(r, *info); auto operandIds = readKnownOperandIds(bc, r, opcode, variant, *info, imms); auto operands = materializeOperands(bc, operandIds); + normalizeLegacyFpOperandOrder(opcode, operands); uint64_t numResults = info->num_results; llvm::SmallVector resultTypes; @@ -620,11 +715,17 @@ static mlir::Operation *buildKnownOpFromReader(BuildCtx &bc, Reader &r, state.addOperands(operands); state.addTypes(resultTypes); addAttrDictionary(state, getAttrDict(bc, attrId)); + auto unifiedOperandSegments = getUnifiedOperandSegments(opcode, operands); addImmediateAttrs(bc, state, *info, imms); for (unsigned i = 0; i < info->num_regions; ++i) (void)state.addRegion(); mlir::Operation *op = mlir::Operation::create(state); + // In MLIR 21, AttrSizedOperandSegments is stored as an inherent property. + // OperationState's generic attribute list does not initialize it, so update + // the generated property storage directly after creating the registered op. + if (unifiedOperandSegments) + setUnifiedOperandSegmentProperty(op, opcode, *unifiedOperandSegments); block.getOperations().push_back(op); registerDecodedOp(bc, opId, op); assignDecodedResults(bc, resStart, op, numResults); diff --git a/tools/ptobc/testdata/fp_operand_forms_v0_roundtrip.pto b/tools/ptobc/testdata/fp_operand_forms_v0_roundtrip.pto new file mode 100644 index 0000000000..c3ef588424 --- /dev/null +++ b/tools/ptobc/testdata/fp_operand_forms_v0_roundtrip.pto @@ -0,0 +1,24 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +module { + func.func private @fp_operand_forms_v0( + %src: !pto.tile_buf, + %fp: !pto.tile_buf, + %row: index, + %col: index, + %dst: !pto.tile_buf) { + pto.textract ins(%src, %row, %col : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + pto.tinsert ins(%src, %row, %col : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + pto.tmov ins(%src : !pto.tile_buf, %fp : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} diff --git a/tools/ptobc/tests/CMakeLists.txt b/tools/ptobc/tests/CMakeLists.txt index 263be1fbee..2afc1f5d67 100644 --- a/tools/ptobc/tests/CMakeLists.txt +++ b/tools/ptobc/tests/CMakeLists.txt @@ -36,6 +36,12 @@ add_test(NAME ptobc_opcode_coverage_check ${CMAKE_SOURCE_DIR}/tools/ptobc/generated/ptobc_opcodes_v0.h ) +add_test(NAME ptobc_v0_fp_schema_compatibility_check + COMMAND ${Python3_EXECUTABLE} + ${CMAKE_CURRENT_LIST_DIR}/v0_fp_schema_compatibility_check.py + ${CMAKE_SOURCE_DIR}/tools/ptobc/generated/ptobc_opcodes_v0.h +) + add_test(NAME ptobc_trowexpandsub_v0_encode COMMAND ${CMAKE_COMMAND} -E env PTOBC_BIN=$ @@ -109,10 +115,21 @@ add_test(NAME ptobc_tconcat_set_validshape_v0_encode add_test(NAME ptobc_tstore_fp_v0_encode COMMAND ${CMAKE_COMMAND} -E env PTOBC_BIN=$ + PTOAS_BIN=${CMAKE_BINARY_DIR}/tools/ptoas/ptoas + PYTHON_EXECUTABLE=${Python3_EXECUTABLE} TESTDATA_DIR=${PTObc_TESTDATA_DIR} ${CMAKE_CURRENT_LIST_DIR}/tstore_fp_v0_encode.sh ) +add_test(NAME ptobc_fp_operand_forms_v0_encode + COMMAND ${CMAKE_COMMAND} -E env + PTOBC_BIN=$ + PTOAS_BIN=${CMAKE_BINARY_DIR}/tools/ptoas/ptoas + PYTHON_EXECUTABLE=${Python3_EXECUTABLE} + TESTDATA_DIR=${PTObc_TESTDATA_DIR} + ${CMAKE_CURRENT_LIST_DIR}/fp_operand_forms_v0_encode.sh +) + add_test(NAME ptobc_comm_p2p_dynamic_v0_encode COMMAND ${CMAKE_COMMAND} -E env PTOBC_BIN=$ diff --git a/tools/ptobc/tests/fp_operand_forms_v0_encode.sh b/tools/ptobc/tests/fp_operand_forms_v0_encode.sh new file mode 100755 index 0000000000..c29d0c98c7 --- /dev/null +++ b/tools/ptobc/tests/fp_operand_forms_v0_encode.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +set -euo pipefail + +PTOBC_BIN=${PTOBC_BIN:-} +PTOAS_BIN=${PTOAS_BIN:-} +TESTDATA_DIR=${TESTDATA_DIR:-} +PYTHON_EXECUTABLE=${PYTHON_EXECUTABLE:-} +if [[ -z "${PTOBC_BIN}" || -z "${PTOAS_BIN}" || -z "${PYTHON_EXECUTABLE}" || -z "${TESTDATA_DIR}" ]]; then + echo "error: PTOBC_BIN, PTOAS_BIN, PYTHON_EXECUTABLE, and TESTDATA_DIR must be set" >&2 + exit 2 +fi + +IN="${TESTDATA_DIR}/fp_operand_forms_v0_roundtrip.pto" +OUT_DIR=${OUT_DIR:-"${PWD}/ptobc_fp_operand_forms_out"} +mkdir -p "${OUT_DIR}" + +BC="${OUT_DIR}/fp_operand_forms_v0_roundtrip.ptobc" +ROUNDTRIP="${OUT_DIR}/fp_operand_forms_v0_roundtrip.roundtrip.pto" + +"${PTOBC_BIN}" encode "${IN}" -o "${BC}" +"${PTOBC_BIN}" decode "${BC}" -o "${ROUNDTRIP}" + +"${PYTHON_EXECUTABLE}" - <<'PY' "${BC}" +from pathlib import Path +import sys + +data = Path(sys.argv[1]).read_bytes() +expected = { + b"\x22\x10": "textract_fp", + b"\x2e\x10": "tinsert_fp", + b"\x39\x10": "tmov.fp", +} +for encoding, name in expected.items(): + if encoding not in data: + raise SystemExit(f"missing legacy {name} opcode encoding") +PY + +grep -F "pto.textract ins(" "${ROUNDTRIP}" >/dev/null +grep -F "pto.tinsert ins(" "${ROUNDTRIP}" >/dev/null +grep -F "pto.tmov ins(" "${ROUNDTRIP}" >/dev/null +[[ $(grep -Fc " fp " "${ROUNDTRIP}") -eq 3 ]] + +# Parsing and verification prove that decoder-reconstructed segment metadata is +# valid, including records encoded with the legacy FP wire operand ordering. +"${PTOAS_BIN}" --pto-arch=a3 --emit-pto-ir "${ROUNDTRIP}" -o /dev/null diff --git a/tools/ptobc/tests/tstore_fp_v0_encode.sh b/tools/ptobc/tests/tstore_fp_v0_encode.sh index 1da36ad102..9098ae4988 100755 --- a/tools/ptobc/tests/tstore_fp_v0_encode.sh +++ b/tools/ptobc/tests/tstore_fp_v0_encode.sh @@ -15,6 +15,18 @@ if [[ -z "${PTOBC_BIN}" ]]; then exit 2 fi +PTOAS_BIN=${PTOAS_BIN:-} +if [[ -z "${PTOAS_BIN}" ]]; then + echo "error: PTOAS_BIN not set" >&2 + exit 2 +fi + +PYTHON_EXECUTABLE=${PYTHON_EXECUTABLE:-} +if [[ -z "${PYTHON_EXECUTABLE}" ]]; then + echo "error: PYTHON_EXECUTABLE not set" >&2 + exit 2 +fi + TESTDATA_DIR=${TESTDATA_DIR:-} if [[ -z "${TESTDATA_DIR}" ]]; then echo "error: TESTDATA_DIR not set" >&2 @@ -31,6 +43,17 @@ ROUNDTRIP="${OUT_DIR}/tstore_fp_v0_roundtrip.roundtrip.pto" "${PTOBC_BIN}" encode "${IN}" -o "${BC}" "${PTOBC_BIN}" decode "${BC}" -o "${ROUNDTRIP}" +"${PYTHON_EXECUTABLE}" - <<'PY' "${BC}" +from pathlib import Path +import sys + +data = Path(sys.argv[1]).read_bytes() +if b"\x66\x10" not in data: + raise SystemExit("missing legacy tstore_fp opcode encoding") +PY + grep -F "pto.tstore ins(" "${ROUNDTRIP}" >/dev/null grep -F " fp " "${ROUNDTRIP}" >/dev/null grep -F "!pto.partition_tensor_view<32x32xi8>" "${ROUNDTRIP}" >/dev/null + +"${PTOAS_BIN}" --pto-arch=a3 --emit-pto-ir "${ROUNDTRIP}" -o /dev/null diff --git a/tools/ptobc/tests/v0_fp_schema_compatibility_check.py b/tools/ptobc/tests/v0_fp_schema_compatibility_check.py new file mode 100755 index 0000000000..e8c5b8d99a --- /dev/null +++ b/tools/ptobc/tests/v0_fp_schema_compatibility_check.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import re +import sys +from pathlib import Path + + +EXPECTED = { + 0x1021: ("pto.textract", 0x00, 4), + 0x1022: ("pto.textract", 0x00, 5), + 0x102D: ("pto.tinsert", 0x00, 4), + 0x102E: ("pto.tinsert", 0x00, 5), + 0x1038: ("pto.tmov", 0x00, 2), + 0x1039: ("pto.tmov", 0x00, 3), + 0x1065: ("pto.tstore", 0x02, 0), + 0x1066: ("pto.tstore", 0x00, 3), +} + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + + text = Path(sys.argv[1]).read_text(encoding="utf-8") + rows = {} + pattern = re.compile( + r'\{0x([0-9A-Fa-f]+),\s*"([^"]+)",\s*\d+,\s*0x[0-9A-Fa-f]+,\s*' + r'0x([0-9A-Fa-f]+),\s*(\d+),' + ) + for opcode, name, operand_mode, num_operands in pattern.findall(text): + rows[int(opcode, 16)] = (name, int(operand_mode, 16), int(num_operands)) + + mismatches = { + f"0x{opcode:04X}": (rows.get(opcode), expected) + for opcode, expected in EXPECTED.items() + if rows.get(opcode) != expected + } + if mismatches: + for opcode, (actual, expected) in mismatches.items(): + print(f"{opcode}: actual={actual}, expected={expected}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 8bf6f45f33692b50412b71501632edf6db25f082 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Tue, 4 Aug 2026 16:19:19 +0800 Subject: [PATCH 005/122] docs: leave release notes unchanged --- ReleaseNotes.md | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/ReleaseNotes.md b/ReleaseNotes.md index 74d7a92b0c..79882f63e0 100644 --- a/ReleaseNotes.md +++ b/ReleaseNotes.md @@ -1,18 +1,5 @@ # PTOAS (PTO Assembler & Optimizer) -## 未发布 - -### PTO IR 接口调整 - -- `pto.tfillpad_inplace` 和 `pto.tfillpad_expand` 已整合到 - `pto.tfillpad`；通过 `mode` 选择行为，默认值为 `normal`。 -- `pto.textract_fp`、`pto.tinsert_fp`、`pto.tmov.fp` 和 - `pto.tstore_fp` 已移除。请改用 `pto.textract`、`pto.tinsert`、 - `pto.tmov` 和 `pto.tstore` 的可选 `fp` operand。旧文本 PTO IR - 不会自动迁移，需要在升级时替换 op 名称并按新语法传入 `fp`。 -- PTO-BC v0 保留上述 FP op 的历史 opcode 作为 wire 兼容别名；已有 - `.ptobc` 文件可由新版 decoder 读取并转换为统一后的 PTO IR。 - ## 版本 - 版本号：v0.51 - 发布日期：2026-02-14 From e488f9e3db10ee3528624860e257c5c1fe70d950 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Tue, 4 Aug 2026 17:17:51 +0800 Subject: [PATCH 006/122] fix: use GitHub pto-isa pin for simulator --- .github/scripts/update_pto_isa_pin.py | 36 --------------------------- .github/workflows/ci_sim.yml | 7 +++--- 2 files changed, 4 insertions(+), 39 deletions(-) diff --git a/.github/scripts/update_pto_isa_pin.py b/.github/scripts/update_pto_isa_pin.py index 81ee204d6c..e9ed2572ce 100755 --- a/.github/scripts/update_pto_isa_pin.py +++ b/.github/scripts/update_pto_isa_pin.py @@ -40,11 +40,6 @@ def parse_args() -> argparse.Namespace: default="docker/Dockerfile", help="Path to the Dockerfile that vendors pto-isa.", ) - parser.add_argument( - "--sim-workflow", - default=".github/workflows/ci_sim.yml", - help="Path to the VPTO simulator workflow file.", - ) parser.add_argument( "--compile-only-guide", default="docs/no_npu_compile_only_guide_zh.md", @@ -150,20 +145,6 @@ def update_dockerfile(path: pathlib.Path, commit: str) -> bool: return False -def update_sim_workflow(path: pathlib.Path, commit: str) -> bool: - original = read_text(path) - updated = replace_exactly_once( - original, - r"^(\s*PTO_ISA_COMMIT:\s*)([0-9a-f]{40})$", - rf"\g<1>{commit}", - path, - ) - if updated != original: - write_text(path, updated) - return True - return False - - def update_compile_only_guide(path: pathlib.Path, commit: str) -> bool: original = read_text(path) updated = replace_exact_count( @@ -222,16 +203,6 @@ def extract_docker_commit(path: pathlib.Path) -> tuple[str, str]: return arg_match.group(1), comment_match.group(1) -def extract_sim_commit(path: pathlib.Path) -> str: - text = read_text(path) - match = re.search( - r"^\s*PTO_ISA_COMMIT:\s*([0-9a-f]{40})$", text, flags=re.MULTILINE - ) - if not match: - raise RuntimeError(f"failed to read pinned pto-isa commit from {path}") - return match.group(1) - - def extract_compile_only_commits(path: pathlib.Path) -> tuple[str, str]: matches = re.findall( r"^export PTO_ISA_COMMIT=([0-9a-f]{40})$", @@ -260,14 +231,12 @@ def extract_remote_validation_commit(path: pathlib.Path) -> str: def verify( ci_path: pathlib.Path, docker_path: pathlib.Path, - sim_workflow_path: pathlib.Path, compile_only_guide_path: pathlib.Path, remote_validation_path: pathlib.Path, commit: str, ) -> None: ci_default, ci_env = extract_ci_commit(ci_path) docker_arg, docker_comment = extract_docker_commit(docker_path) - sim_commit = extract_sim_commit(sim_workflow_path) guide_setup, guide_run = extract_compile_only_commits(compile_only_guide_path) remote_validation_commit = extract_remote_validation_commit( remote_validation_path @@ -277,7 +246,6 @@ def verify( f"{ci_path}:runtime_default": ci_env, f"{docker_path}:arg": docker_arg, f"{docker_path}:comment": docker_comment, - f"{sim_workflow_path}:simulator": sim_commit, f"{compile_only_guide_path}:setup": guide_setup, f"{compile_only_guide_path}:run": guide_run, f"{remote_validation_path}:fallback": remote_validation_commit, @@ -293,7 +261,6 @@ def main() -> int: commit = args.commit or resolve_head_commit(args.repo_url) ci_path = pathlib.Path(args.ci_workflow) docker_path = pathlib.Path(args.dockerfile) - sim_workflow_path = pathlib.Path(args.sim_workflow) compile_only_guide_path = pathlib.Path(args.compile_only_guide) remote_validation_path = pathlib.Path(args.remote_validation_script) @@ -301,7 +268,6 @@ def main() -> int: verify( ci_path, docker_path, - sim_workflow_path, compile_only_guide_path, remote_validation_path, commit, @@ -311,13 +277,11 @@ def main() -> int: update_ci_workflow(ci_path, commit) update_dockerfile(docker_path, commit) - update_sim_workflow(sim_workflow_path, commit) update_compile_only_guide(compile_only_guide_path, commit) update_remote_validation_script(remote_validation_path, commit) verify( ci_path, docker_path, - sim_workflow_path, compile_only_guide_path, remote_validation_path, commit, diff --git a/.github/workflows/ci_sim.yml b/.github/workflows/ci_sim.yml index 65b59983b9..011c1bdcd1 100644 --- a/.github/workflows/ci_sim.yml +++ b/.github/workflows/ci_sim.yml @@ -116,9 +116,10 @@ jobs: PYPTO_REF: ef6ce7cd8bd33b4c93b58dc34830a20b145736ac PYPTO_WORKSPACE: ${{ github.workspace }}/.work/pypto-ci PYPTO_RUN_WORKSPACE: ${{ github.workspace }}/.work/pypto-run-ci - # Keep this pin aligned with the compiler, Docker, documentation, and - # remote-validation defaults via update_pto_isa_pin.py. - PTO_ISA_COMMIT: 27386d906e8fdcbd93aec84197939bc0b2c6caea + # GitHub pin containing the unified TFILLPAD API, the two-argument Soft + # SYNCALL ABI, and the follow-up CPU-Sim duplicate-stub fix. Keep this + # separate from the GitCode pins managed by update_pto_isa_pin.py. + PTO_ISA_COMMIT: e948507e18ec4f39037a04914b97e77f5b9d75e3 PTO_ISA_ROOT: ${{ github.workspace }}/.work/pto-isa-ci steps: - name: Checkout From 162ddbca7af418a7df4f052f641e63b7a24348a1 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Tue, 4 Aug 2026 19:11:19 +0800 Subject: [PATCH 007/122] fix: use pinned pto-isa in tilelang st --- test/lit/pto/tinsert_a5_extended_modes.pto | 2 +- test/tilelang_st/npu/a5/src/st/CMakeLists.txt | 6 ++++++ test/tilelang_st/npu/a5/src/st/smoke/CMakeLists.txt | 6 ++++++ test/tilelang_st/script/run_st.py | 5 ++++- 4 files changed, 17 insertions(+), 2 deletions(-) diff --git a/test/lit/pto/tinsert_a5_extended_modes.pto b/test/lit/pto/tinsert_a5_extended_modes.pto index d20411c647..6390abf26c 100644 --- a/test/lit/pto/tinsert_a5_extended_modes.pto +++ b/test/lit/pto/tinsert_a5_extended_modes.pto @@ -45,7 +45,7 @@ module attributes {"pto.target_arch" = "a5"} { // A5: acc->mat with fp (vector) quantization (generic MLIR form). // CHECK-LABEL: AICORE void tinsert_acc_mat_fp_quant( - // CHECK: TINSERT<{{.*}}TileType::Mat, int8_t{{.*}}, {{.*}}TileType::Acc, float{{.*}}, {{.*}}TileType::Scaling{{.*}}>( + // CHECK: TINSERT_FP( func.func @tinsert_acc_mat_fp_quant() { %c0 = arith.constant 0 : index %src = pto.alloc_tile : !pto.tile_buf diff --git a/test/tilelang_st/npu/a5/src/st/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/CMakeLists.txt index a3e167a846..0cdc9b8524 100644 --- a/test/tilelang_st/npu/a5/src/st/CMakeLists.txt +++ b/test/tilelang_st/npu/a5/src/st/CMakeLists.txt @@ -40,6 +40,9 @@ else() endif() set(PTO_ISA_ROOT "${CMAKE_CURRENT_LIST_DIR}/../../../../../../../pto-isa" CACHE PATH "Path to pto-isa repo") +if(NOT EXISTS "${PTO_ISA_ROOT}/include/pto/pto-inst.hpp") + message(FATAL_ERROR "Cannot find ${PTO_ISA_ROOT}/include/pto/pto-inst.hpp") +endif() set(PTO_TILELANG_ST_COMMON_DIR "${CMAKE_CURRENT_LIST_DIR}/common") set(ASCEND_DRIVER_PATH /usr/local/Ascend/driver) @@ -82,5 +85,8 @@ include_directories( ${ASCEND_HOME_PATH}/include ${ASCEND_DRIVER_PATH}/kernel/inc ) +include_directories(BEFORE + ${PTO_ISA_ROOT}/include +) add_subdirectory(testcase) diff --git a/test/tilelang_st/npu/a5/src/st/smoke/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/smoke/CMakeLists.txt index 66f2852583..4ce1bf87d2 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/CMakeLists.txt +++ b/test/tilelang_st/npu/a5/src/st/smoke/CMakeLists.txt @@ -40,6 +40,9 @@ else() endif() set(PTO_ISA_ROOT "${CMAKE_CURRENT_LIST_DIR}/../../../../../../../../pto-isa" CACHE PATH "Path to pto-isa repo") +if(NOT EXISTS "${PTO_ISA_ROOT}/include/pto/pto-inst.hpp") + message(FATAL_ERROR "Cannot find ${PTO_ISA_ROOT}/include/pto/pto-inst.hpp") +endif() set(PTO_TILELANG_ST_COMMON_DIR "${CMAKE_CURRENT_LIST_DIR}/../common") set(ASCEND_DRIVER_PATH /usr/local/Ascend/driver) @@ -82,5 +85,8 @@ include_directories( ${ASCEND_HOME_PATH}/include ${ASCEND_DRIVER_PATH}/kernel/inc ) +include_directories(BEFORE + ${PTO_ISA_ROOT}/include +) add_subdirectory(testcase) diff --git a/test/tilelang_st/script/run_st.py b/test/tilelang_st/script/run_st.py index db5fe52eff..f03a6d75f9 100755 --- a/test/tilelang_st/script/run_st.py +++ b/test/tilelang_st/script/run_st.py @@ -129,8 +129,11 @@ def build_project(run_mode, soc_version, testcase, ptoas_bin, build_jobs=None): f"-DSOC_VERSION={soc_version}", f"-DTEST_CASE={testcase}", f"-DPTOAS_BIN={ptoas_bin}", - "..", ] + pto_isa_root = os.environ.get("PTO_ISA_ROOT") + if pto_isa_root: + cmake_cmd.append(f"-DPTO_ISA_ROOT={os.path.abspath(pto_isa_root)}") + cmake_cmd.append("..") subprocess.run( cmake_cmd, cwd=build_dir, From 3308e0649e30f77a5d7e834a7ac117f41307cf96 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Tue, 4 Aug 2026 19:53:46 +0800 Subject: [PATCH 008/122] test: exercise actual tfillpad inplace padding --- .../smoke/testcase/tfillpad_inplace/cases.py | 8 +-- .../testcase/tfillpad_inplace/gen_data.py | 9 ++- .../smoke/testcase/tfillpad_inplace/main.cpp | 4 +- .../tfillpad_inplace/tfillpad_inplace.pto | 61 ++++++++++--------- .../src/st/testcase/tfillpad_inplace/cases.py | 8 +-- .../st/testcase/tfillpad_inplace/gen_data.py | 11 ++-- .../src/st/testcase/tfillpad_inplace/main.cpp | 6 +- .../tfillpad_inplace/tfillpad_inplace.pto | 37 ++++++----- 8 files changed, 76 insertions(+), 68 deletions(-) diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/cases.py b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/cases.py index 60e45f45c0..72924bac42 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/cases.py +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/cases.py @@ -27,22 +27,22 @@ "name": "f32_64x16_noexpand", "dtype": np.float32, "src_shape": (64, 16), - "src_valid": (64, 16), + "src_valid": (64, 7), "dst_shape": (64, 16), "dst_valid": (64, 16), "fill_padval": "Max", "eps": 1e-6, }, - # ========== Case: float, src_valid == dst_valid (no expansion) ========== + # ========== Case: float, fill cols 7..15 in place ========== { "name": "f32_260x16_noexpand", "dtype": np.float32, "src_shape": (260, 16), # src physical - "src_valid": (260, 16), # src valid = dst valid (no expansion) + "src_valid": (260, 7), # valid data before in-place padding "dst_shape": (260, 16), # dst physical "dst_valid": (260, 16), # dst valid = full output - "fill_padval": "Max", # FillPadVal (not used since no expansion) + "fill_padval": "Max", # fill cols 7..15 in the same tile "eps": 1e-6, }, ] diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/gen_data.py b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/gen_data.py index 5dc09477af..db085a5a99 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/gen_data.py +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/gen_data.py @@ -73,15 +73,18 @@ def save_case_data(case_name, data_dict): dst_r, dst_c = dst_shape dst_vr, dst_vc = dst_valid - # Input: src valid region data (random values) - input_data = np.random.uniform(1.0, 10.0, size=(src_vr, src_vc)).astype(dtype) + # Input occupies the full physical buffer; only src_valid contains data. + input_data = np.zeros(src_shape, dtype=dtype) + input_data[:src_vr, :src_vc] = np.random.uniform( + 1.0, 10.0, size=(src_vr, src_vc) + ).astype(dtype) # Golden: dst full region # Copy src.valid region to dst[:src_vr, :src_vc] # Fill cols src_vc to dst_vc with FillPadVal # Fill rows src_vr to dst_vr with FillPadVal (row expansion, if any) golden = np.zeros(dst_shape, dtype=dtype) - golden[:src_vr, :src_vc] = input_data + golden[:src_vr, :src_vc] = input_data[:src_vr, :src_vc] # Fill column padding (cols src_vc to dst_vc) if dst_vc > src_vc: diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/main.cpp b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/main.cpp index d429171864..1ab11eeb1e 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/main.cpp @@ -41,10 +41,10 @@ struct TestCase { static const TestCase kCases[] = { {"f32_64x16_noexpand", DataType::F32, LaunchTFILLPAD_INPLACE_f32_64x16_noexpand, - 64, 16, 64, 16, sizeof(float)}, + 64, 16, 64, 7, sizeof(float)}, {"f32_260x16_noexpand", DataType::F32, LaunchTFILLPAD_INPLACE_f32_260x16_noexpand, - 260, 16, 260, 16, sizeof(float)}, + 260, 16, 260, 7, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto index 4421609d6d..0529c2ab8b 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto @@ -13,18 +13,17 @@ // PadValue encoding: 0=Null, 1=Zero, 2=Max, 3=Min // Case 5: float, 260x16, valid=260x7, FillPad=Max (pad=2) // -// Note: PTOAS tstore requires dst size to match src valid_shape. -// For outputting full buffer after inplace fill, we use two tiles: -// - src tile: holds input data (valid=260x7) -// - dst tile: receives filled data (valid=260x16 for output) +// After filling the physical padding in place, set the runtime valid shape to +// the full tile extent so tstore writes the filled columns too. module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { - // ========== Smoke case: float, 64x16 physical, src_valid == dst_valid ========== + // ========== Smoke case: float, 64x16 physical, src valid 64x7 ========== func.func @TFILLPAD_INPLACE_f32_64x16_noexpand(%tile_ptr: !pto.ptr) attributes {pto.kernel} { %c0 = arith.constant 0 : index %c0_i64 = arith.constant 0 : i64 %c1 = arith.constant 1 : index + %c7 = arith.constant 7 : index %c16 = arith.constant 16 : index %c64 = arith.constant 64 : index %c1024 = arith.constant 1024 : index @@ -41,35 +40,39 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind -> !pto.partition_tensor_view<1x1x1x64x16xf32> + sizes = [%c1, %c1, %c1, %c64, %c7] + : !pto.tensor_view<1x1x1x64x16xf32> -> !pto.partition_tensor_view<1x1x1x64x7xf32> %dst_part = pto.partition_view %dst_view, offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c64, %c16] : !pto.tensor_view<1x1x1x64x16xf32> -> !pto.partition_tensor_view<1x1x1x64x16xf32> - %tile_buf = pto.alloc_tile addr = %c0_i64 - : !pto.tile_buf + %tile_buf = pto.alloc_tile addr = %c0_i64 valid_row = %c64 valid_col = %c7 + : !pto.tile_buf - pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x64x16xf32>) - outs(%tile_buf : !pto.tile_buf) + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x64x7xf32>) + outs(%tile_buf : !pto.tile_buf) - pto.tfillpad ins(%tile_buf : !pto.tile_buf) - outs(%tile_buf : !pto.tile_buf) + pto.tfillpad ins(%tile_buf : !pto.tile_buf) + outs(%tile_buf : !pto.tile_buf) {mode = #pto.tfillpad_mode} - pto.tstore ins(%tile_buf : !pto.tile_buf) + pto.set_validshape %tile_buf, %c64, %c16 + : !pto.tile_buf + + pto.tstore ins(%tile_buf : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x64x16xf32>) return } - // ========== No expansion: float, 260x16 physical, src_valid == dst_valid ========== + // ========== Float, 260x16 physical, src valid 260x7 ========== func.func @TFILLPAD_INPLACE_f32_260x16_noexpand(%tile_ptr: !pto.ptr) attributes {pto.kernel} { %c0 = arith.constant 0 : index %c0_i64 = arith.constant 0 : i64 %c1 = arith.constant 1 : index + %c7 = arith.constant 7 : index %c16 = arith.constant 16 : index %c260 = arith.constant 260 : index %c4160 = arith.constant 4160 : index // 260*16 (full tile size) @@ -88,29 +91,29 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind -> !pto.partition_tensor_view<1x1x1x260x16xf32> + sizes = [%c1, %c1, %c1, %c260, %c7] + : !pto.tensor_view<1x1x1x260x16xf32> -> !pto.partition_tensor_view<1x1x1x260x7xf32> %dst_part = pto.partition_view %dst_view, offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c260, %c16] : !pto.tensor_view<1x1x1x260x16xf32> -> !pto.partition_tensor_view<1x1x1x260x16xf32> - // Single tile buffer in UB space at address 0 - // src_valid = dst_valid = 260x16, so no expansion needed - %tile_buf = pto.alloc_tile addr = %c0_i64 - : !pto.tile_buf + // Single tile buffer in UB space at address 0. + %tile_buf = pto.alloc_tile addr = %c0_i64 valid_row = %c260 valid_col = %c7 + : !pto.tile_buf - // Load full tile (260x16) - pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x260x16xf32>) - outs(%tile_buf : !pto.tile_buf) + // Load only the valid 260x7 region. + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x260x7xf32>) + outs(%tile_buf : !pto.tile_buf) - // tfillpad in_place: src_valid == dst_valid, no expansion - pto.tfillpad ins(%tile_buf : !pto.tile_buf) - outs(%tile_buf : !pto.tile_buf) + pto.tfillpad ins(%tile_buf : !pto.tile_buf) + outs(%tile_buf : !pto.tile_buf) {mode = #pto.tfillpad_mode} - // Store full tile - pto.tstore ins(%tile_buf : !pto.tile_buf) + pto.set_validshape %tile_buf, %c260, %c16 + : !pto.tile_buf + + pto.tstore ins(%tile_buf : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x16xf32>) return } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/cases.py index c1b1dae177..27fdf40dca 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/cases.py @@ -23,16 +23,16 @@ import numpy as np CASES = [ - # ========== Case: float, src_valid == dst_valid (no expansion) ========== + # ========== Case: float, fill cols 7..15 in place ========== { "name": "f32_260x16_noexpand", "dtype": np.float32, "src_shape": (260, 16), # src physical - "src_valid": (260, 16), # src valid = dst valid (no expansion) + "src_valid": (260, 7), # valid data before in-place padding "dst_shape": (260, 16), # dst physical "dst_valid": (260, 16), # dst valid = full output - "fill_padval": "Max", # FillPadVal (not used since no expansion) + "fill_padval": "Max", # fill cols 7..15 in the same tile "eps": 1e-6, }, -] \ No newline at end of file +] diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/gen_data.py index 2345a38a4c..db085a5a99 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/gen_data.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/gen_data.py @@ -73,15 +73,18 @@ def save_case_data(case_name, data_dict): dst_r, dst_c = dst_shape dst_vr, dst_vc = dst_valid - # Input: src valid region data (random values) - input_data = np.random.uniform(1.0, 10.0, size=(src_vr, src_vc)).astype(dtype) + # Input occupies the full physical buffer; only src_valid contains data. + input_data = np.zeros(src_shape, dtype=dtype) + input_data[:src_vr, :src_vc] = np.random.uniform( + 1.0, 10.0, size=(src_vr, src_vc) + ).astype(dtype) # Golden: dst full region # Copy src.valid region to dst[:src_vr, :src_vc] # Fill cols src_vc to dst_vc with FillPadVal # Fill rows src_vr to dst_vr with FillPadVal (row expansion, if any) golden = np.zeros(dst_shape, dtype=dtype) - golden[:src_vr, :src_vc] = input_data + golden[:src_vr, :src_vc] = input_data[:src_vr, :src_vc] # Fill column padding (cols src_vc to dst_vc) if dst_vc > src_vc: @@ -96,4 +99,4 @@ def save_case_data(case_name, data_dict): save_case_data(case["name"], {"input": input_data, "golden": golden}) print(f"[INFO] gen_data: {case['name']} " f"src_valid={src_valid} dst_shape={dst_shape} " - f"fill_pad={fill_padval} dtype={dtype.__name__}") \ No newline at end of file + f"fill_pad={fill_padval} dtype={dtype.__name__}") diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/main.cpp index 34f94bf408..7ab5270618 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/main.cpp @@ -37,9 +37,9 @@ struct TestCase { }; static const TestCase kCases[] = { - // Case: float, 260x16, no expansion (inplace: single buffer) + // Case: float, fill cols 7..15 in place in a single buffer. {"f32_260x16_noexpand", DataType::F32, - 260, 16, 260, 16, sizeof(float)}, + 260, 16, 260, 7, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); @@ -126,4 +126,4 @@ int main(int argc, char *argv[]) { aclFinalize(); return rc; -} \ No newline at end of file +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto index 36cf1dff95..0f61dab617 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto @@ -13,18 +13,17 @@ // PadValue encoding: 0=Null, 1=Zero, 2=Max, 3=Min // Case 5: float, 260x16, valid=260x7, FillPad=Max (pad=2) // -// Note: PTOAS tstore requires dst size to match src valid_shape. -// For outputting full buffer after inplace fill, we use two tiles: -// - src tile: holds input data (valid=260x7) -// - dst tile: receives filled data (valid=260x16 for output) +// After filling the physical padding in place, set the runtime valid shape to +// the full tile extent so tstore writes the filled columns too. module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { - // ========== No expansion: float, 260x16 physical, src_valid == dst_valid ========== + // ========== Float, 260x16 physical, src valid 260x7 ========== func.func @TFILLPAD_INPLACE_f32_260x16_noexpand(%tile_ptr: !pto.ptr) attributes {pto.kernel} { %c0 = arith.constant 0 : index %c0_i64 = arith.constant 0 : i64 %c1 = arith.constant 1 : index + %c7 = arith.constant 7 : index %c16 = arith.constant 16 : index %c260 = arith.constant 260 : index %c4160 = arith.constant 4160 : index // 260*16 (full tile size) @@ -43,29 +42,29 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind -> !pto.partition_tensor_view<1x1x1x260x16xf32> + sizes = [%c1, %c1, %c1, %c260, %c7] + : !pto.tensor_view<1x1x1x260x16xf32> -> !pto.partition_tensor_view<1x1x1x260x7xf32> %dst_part = pto.partition_view %dst_view, offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c260, %c16] : !pto.tensor_view<1x1x1x260x16xf32> -> !pto.partition_tensor_view<1x1x1x260x16xf32> - // Single tile buffer in UB space at address 0 - // src_valid = dst_valid = 260x16, so no expansion needed - %tile_buf = pto.alloc_tile addr = %c0_i64 - : !pto.tile_buf + // Single tile buffer in UB space at address 0. + %tile_buf = pto.alloc_tile addr = %c0_i64 valid_row = %c260 valid_col = %c7 + : !pto.tile_buf - // Load full tile (260x16) - pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x260x16xf32>) - outs(%tile_buf : !pto.tile_buf) + // Load only the valid 260x7 region. + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x260x7xf32>) + outs(%tile_buf : !pto.tile_buf) - // tfillpad in_place: src_valid == dst_valid, no expansion - pto.tfillpad ins(%tile_buf : !pto.tile_buf) - outs(%tile_buf : !pto.tile_buf) + pto.tfillpad ins(%tile_buf : !pto.tile_buf) + outs(%tile_buf : !pto.tile_buf) {mode = #pto.tfillpad_mode} - // Store full tile - pto.tstore ins(%tile_buf : !pto.tile_buf) + pto.set_validshape %tile_buf, %c260, %c16 + : !pto.tile_buf + + pto.tstore ins(%tile_buf : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x16xf32>) return } From d012398f38814ce5f9ac47e4bc80958e8482bee5 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Tue, 4 Aug 2026 20:26:00 +0800 Subject: [PATCH 009/122] test: match pto-isa tfillpad inplace aliasing --- .../testcase/tfillpad_inplace/gen_data.py | 2 +- .../tfillpad_inplace/tfillpad_inplace.pto | 41 +++++++++---------- .../st/testcase/tfillpad_inplace/gen_data.py | 2 +- .../tfillpad_inplace/tfillpad_inplace.pto | 25 +++++------ 4 files changed, 35 insertions(+), 35 deletions(-) diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/gen_data.py b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/gen_data.py index db085a5a99..4acc371531 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/gen_data.py +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/gen_data.py @@ -11,7 +11,7 @@ """Generate golden data for tfillpad_inplace test cases. For tfillpad_inplace: - - Only one tile, valid_shape smaller than tile shape + - Source and destination tile handles alias the same UB storage - Input: full tile shape (rows x cols), random values in valid region, zeros in padding - Golden: full tile shape with valid region copied and padding filled with MAX (PadValue.Max) """ diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto index 0529c2ab8b..1c2b40183b 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto @@ -13,8 +13,9 @@ // PadValue encoding: 0=Null, 1=Zero, 2=Max, 3=Min // Case 5: float, 260x16, valid=260x7, FillPad=Max (pad=2) // -// After filling the physical padding in place, set the runtime valid shape to -// the full tile extent so tstore writes the filled columns too. +// Source and destination tile handles share UB address 0, matching the +// PTO-ISA Case 5 calling convention. Source handles keep the x7 valid shape; +// destination handles expose the full shape consumed by tstore. module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { // ========== Smoke case: float, 64x16 physical, src valid 64x7 ========== @@ -47,20 +48,19 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind -> !pto.partition_tensor_view<1x1x1x64x16xf32> - %tile_buf = pto.alloc_tile addr = %c0_i64 valid_row = %c64 valid_col = %c7 - : !pto.tile_buf + %src_tile = pto.alloc_tile addr = %c0_i64 valid_row = %c64 valid_col = %c7 + : !pto.tile_buf + %dst_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x64x7xf32>) - outs(%tile_buf : !pto.tile_buf) + outs(%src_tile : !pto.tile_buf) - pto.tfillpad ins(%tile_buf : !pto.tile_buf) - outs(%tile_buf : !pto.tile_buf) + pto.tfillpad ins(%src_tile : !pto.tile_buf) + outs(%dst_tile : !pto.tile_buf) {mode = #pto.tfillpad_mode} - pto.set_validshape %tile_buf, %c64, %c16 - : !pto.tile_buf - - pto.tstore ins(%tile_buf : !pto.tile_buf) + pto.tstore ins(%dst_tile : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x64x16xf32>) return } @@ -98,22 +98,21 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind -> !pto.partition_tensor_view<1x1x1x260x16xf32> - // Single tile buffer in UB space at address 0. - %tile_buf = pto.alloc_tile addr = %c0_i64 valid_row = %c260 valid_col = %c7 - : !pto.tile_buf + // Distinct tile handles alias the same UB storage. + %src_tile = pto.alloc_tile addr = %c0_i64 valid_row = %c260 valid_col = %c7 + : !pto.tile_buf + %dst_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf // Load only the valid 260x7 region. pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x260x7xf32>) - outs(%tile_buf : !pto.tile_buf) + outs(%src_tile : !pto.tile_buf) - pto.tfillpad ins(%tile_buf : !pto.tile_buf) - outs(%tile_buf : !pto.tile_buf) + pto.tfillpad ins(%src_tile : !pto.tile_buf) + outs(%dst_tile : !pto.tile_buf) {mode = #pto.tfillpad_mode} - pto.set_validshape %tile_buf, %c260, %c16 - : !pto.tile_buf - - pto.tstore ins(%tile_buf : !pto.tile_buf) + pto.tstore ins(%dst_tile : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x16xf32>) return } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/gen_data.py index db085a5a99..4acc371531 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/gen_data.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/gen_data.py @@ -11,7 +11,7 @@ """Generate golden data for tfillpad_inplace test cases. For tfillpad_inplace: - - Only one tile, valid_shape smaller than tile shape + - Source and destination tile handles alias the same UB storage - Input: full tile shape (rows x cols), random values in valid region, zeros in padding - Golden: full tile shape with valid region copied and padding filled with MAX (PadValue.Max) """ diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto index 0f61dab617..32e797a038 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto @@ -13,8 +13,9 @@ // PadValue encoding: 0=Null, 1=Zero, 2=Max, 3=Min // Case 5: float, 260x16, valid=260x7, FillPad=Max (pad=2) // -// After filling the physical padding in place, set the runtime valid shape to -// the full tile extent so tstore writes the filled columns too. +// The source and destination tile handles share UB address 0, matching the +// PTO-ISA Case 5 calling convention. The source keeps the 260x7 valid shape; +// the destination has the full 260x16 shape consumed by tstore. module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { // ========== Float, 260x16 physical, src valid 260x7 ========== @@ -49,22 +50,22 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind -> !pto.partition_tensor_view<1x1x1x260x16xf32> - // Single tile buffer in UB space at address 0. - %tile_buf = pto.alloc_tile addr = %c0_i64 valid_row = %c260 valid_col = %c7 - : !pto.tile_buf + // Distinct tile handles alias the same UB storage. This mirrors + // vecTile/vecTileP in the PTO-ISA reference implementation. + %src_tile = pto.alloc_tile addr = %c0_i64 valid_row = %c260 valid_col = %c7 + : !pto.tile_buf + %dst_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf // Load only the valid 260x7 region. pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x260x7xf32>) - outs(%tile_buf : !pto.tile_buf) + outs(%src_tile : !pto.tile_buf) - pto.tfillpad ins(%tile_buf : !pto.tile_buf) - outs(%tile_buf : !pto.tile_buf) + pto.tfillpad ins(%src_tile : !pto.tile_buf) + outs(%dst_tile : !pto.tile_buf) {mode = #pto.tfillpad_mode} - pto.set_validshape %tile_buf, %c260, %c16 - : !pto.tile_buf - - pto.tstore ins(%tile_buf : !pto.tile_buf) + pto.tstore ins(%dst_tile : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x16xf32>) return } From 33d66acc668c0e718d2aa14f349173c97930582b Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Wed, 5 Aug 2026 00:22:39 +0800 Subject: [PATCH 010/122] fix(vpto): add SIMT fastmath control --- test/lit/vpto/bisheng_simt_fastmath_cli.pto | 14 ++++++++++++++ tools/ptoas/ObjectEmission.cpp | 9 +++++++++ 2 files changed, 23 insertions(+) create mode 100644 test/lit/vpto/bisheng_simt_fastmath_cli.pto diff --git a/test/lit/vpto/bisheng_simt_fastmath_cli.pto b/test/lit/vpto/bisheng_simt_fastmath_cli.pto new file mode 100644 index 0000000000..4fc4dbd27f --- /dev/null +++ b/test/lit/vpto/bisheng_simt_fastmath_cli.pto @@ -0,0 +1,14 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software; you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --help | FileCheck %s +// RUN: ptoas --simt-fastmath --help > /dev/null +// RUN: ptoas --simt-fastmath=false --help > /dev/null + +// CHECK: --simt-fastmath +// CHECK-SAME: Enable Bisheng SIMT floating-point contraction and fast math combining for VPTO device compilation diff --git a/tools/ptoas/ObjectEmission.cpp b/tools/ptoas/ObjectEmission.cpp index 806b0cfc9f..5ca4567aca 100644 --- a/tools/ptoas/ObjectEmission.cpp +++ b/tools/ptoas/ObjectEmission.cpp @@ -50,6 +50,12 @@ static llvm::cl::opt enableBishengVecMISched( "the scheduler"), llvm::cl::init(false)); +static llvm::cl::opt enableSimtFastMath( + "simt-fastmath", + llvm::cl::desc("Enable Bisheng SIMT floating-point contraction and fast " + "math combining for VPTO device compilation"), + llvm::cl::init(false)); + static llvm::cl::opt bishengVFAutoSyncMode( "bisheng-vf-auto-sync", llvm::cl::desc("Explicit Bisheng VF auto-sync mode for VPTO device " @@ -520,6 +526,9 @@ static bool compileDeviceLLVMToObject(llvm::StringRef llPath, args.push_back("-mllvm"); args.push_back("--cce-aicore-vec-misched=0"); } + args.push_back("-mllvm"); + args.push_back(std::string("--cce-simt-fpmath-combine=") + + (enableSimtFastMath ? "true" : "false")); args.push_back("-c"); args.push_back("-x"); args.push_back("ir"); From c090ce20456cb5a6ec8d97c63769e78695d7a216 Mon Sep 17 00:00:00 2001 From: andodo Date: Wed, 5 Aug 2026 09:34:23 +0800 Subject: [PATCH 011/122] feat(ptodsl): expose explicit L1 to L0 loads --- .../docs/user_guide/07-data-movement-ops.md | 28 ++++++++++ ptodsl/ptodsl/_ops.py | 54 +++++++++++++++++++ ptodsl/ptodsl/pto.py | 1 + ptodsl/tests/test_jit_compile.py | 35 ++++++++++++ 4 files changed, 118 insertions(+) diff --git a/ptodsl/docs/user_guide/07-data-movement-ops.md b/ptodsl/docs/user_guide/07-data-movement-ops.md index 76557fdf7d..042d3474e5 100644 --- a/ptodsl/docs/user_guide/07-data-movement-ops.md +++ b/ptodsl/docs/user_guide/07-data-movement-ops.md @@ -939,6 +939,32 @@ Cube compute step; it does not issue those transfers itself. ### Operand loading: L1 → L0A / L0B +#### `pto.load_cbuf_to_ca(src: PtrType, dst: PtrType, m_start: int, k_start: int, m_step: int, k_step: int, src_stride: int, dst_stride: int, *, transpose: bool = False) -> None` +#### `pto.load_cbuf_to_cb(src: PtrType, dst: PtrType, m_start: int, k_start: int, m_step: int, k_step: int, src_stride: int, dst_stride: int, *, transpose: bool = False) -> None` + +**Description**: Explicit-control L1-to-L0A/L0B loads. Unlike the structured +`mte_l1_l0a/b` wrappers, these APIs preserve the authored fractal-block control +fields and do not infer strides from a logical tile shape. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `src` | `PtrType` (L1/MAT) | L1 source pointer; parent-allocation matrix offsets are represented by the control fields below. | +| `dst` | `PtrType` (L0A/L0B) | L0 destination pointer; stage/version offsets belong in this pointer. | +| `m_start`, `k_start` | `int` | Source fractal-block coordinates at which the load begins. | +| `m_step`, `k_step` | `int` | Number of fractal blocks transferred along the two source axes. | +| `src_stride`, `dst_stride` | `int` | Physical outer strides of the complete L1 and L0 allocations, in fractal-block units. They are independent of the transferred region extents. | +| `transpose` | `bool` | Final hardware transpose attribute. | + +**Returns**: None (side-effect operation). + +Use these APIs when a source/destination subregion has layout control that cannot +be reconstructed from a canonical `m`/`k`/`n` tile shape. The pointer order is +always `src, dst`. + +--- + #### `pto.mte_l1_l0a(src: PtrType, dst: PtrType, m: int, k: int, *, start_row: int, start_col: int, transpose: bool = False) -> None` **Description**: Structured L1-to-L0A (left-operand buffer) load. @@ -1172,6 +1198,8 @@ pto.mte_l0c_gm( | GM → L1 | `mte_gm_l1` | gm | l1 | | GM → L1 (fractal) | `mte_gm_l1_frac` | gm | l1 | | L1 → UB | `mte_l1_ub` | l1 | ub | +| L1 → L0A (explicit control) | `load_cbuf_to_ca` | l1 | l0a | +| L1 → L0B (explicit control) | `load_cbuf_to_cb` | l1 | l0b | | L1 → L0A | `mte_l1_l0a` | l1 | l0a | | L1 → L0B | `mte_l1_l0b` | l1 | l0b | | L1 → L0A (MX) | `mte_l1_l0a_mx` | l1 | l0a | diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 5b0a8298dc..35c3c152fd 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -5276,6 +5276,60 @@ def mem_bar(barrier_type): _pto.MemBarOp(kind=_membar_attr(barrier_name)) +@_explicit_mode_only("pto.load_cbuf_to_ca(...)") +def load_cbuf_to_ca( + source, + destination, + m_start, + k_start, + m_step, + k_step, + src_stride, + dst_stride, + *, + transpose=False, +): + """``pto.load_cbuf_to_ca`` – explicit-control L1-to-L0A load.""" + _pto.LoadCbufToCaOp( + unwrap_surface_value(source), + unwrap_surface_value(destination), + _coerce_i64(m_start, context="load_cbuf_to_ca m_start"), + _coerce_i64(k_start, context="load_cbuf_to_ca k_start"), + _coerce_i64(m_step, context="load_cbuf_to_ca m_step"), + _coerce_i64(k_step, context="load_cbuf_to_ca k_step"), + _coerce_i64(src_stride, context="load_cbuf_to_ca src_stride"), + _coerce_i64(dst_stride, context="load_cbuf_to_ca dst_stride"), + transpose=transpose, + ) + + +@_explicit_mode_only("pto.load_cbuf_to_cb(...)") +def load_cbuf_to_cb( + source, + destination, + m_start, + k_start, + m_step, + k_step, + src_stride, + dst_stride, + *, + transpose=False, +): + """``pto.load_cbuf_to_cb`` – explicit-control L1-to-L0B load.""" + _pto.LoadCbufToCbOp( + unwrap_surface_value(source), + unwrap_surface_value(destination), + _coerce_i64(m_start, context="load_cbuf_to_cb m_start"), + _coerce_i64(k_start, context="load_cbuf_to_cb k_start"), + _coerce_i64(m_step, context="load_cbuf_to_cb m_step"), + _coerce_i64(k_step, context="load_cbuf_to_cb k_step"), + _coerce_i64(src_stride, context="load_cbuf_to_cb src_stride"), + _coerce_i64(dst_stride, context="load_cbuf_to_cb dst_stride"), + transpose=transpose, + ) + + @_explicit_mode_only("pto.mte_l1_l0a(...)") def mte_l1_l0a( source, diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index 5cb8b96168..27418a04b4 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -127,6 +127,7 @@ set_atomic_add, set_atomic_max, set_atomic_min, set_atomic_none, set_atomic_f32, set_atomic_f16, set_atomic_bf16, set_atomic_s32, set_atomic_s16, set_atomic_s8, + load_cbuf_to_ca, load_cbuf_to_cb, mte_l1_l0a, mte_l1_l0b, mte_l1_l0a_mx, mte_l1_l0b_mx, mte_l0c_l1, mte_l0c_gm, mte_l0c_ub, mad, mad_acc, mad_bias, mad_mx, mad_mx_acc, mad_mx_bias, diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 45dc7975f8..a962a135ea 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -2391,6 +2391,28 @@ def public_surface_exports_probe( memory_space=pto.MemorySpace.RIGHT, valid_shape=[16, 16], ) + pto.load_cbuf_to_ca( + lhs_tile.as_ptr(), + lhs_l0a.as_ptr(), + 4, + 0, + 4, + 16, + 16, + 4, + transpose=True, + ) + pto.load_cbuf_to_cb( + rhs_tile.as_ptr(), + rhs_l0b.as_ptr(), + 4, + 0, + 4, + 16, + 16, + 4, + transpose=True, + ) lhs_l0a_mx = pto.alloc_tile( shape=[16, 32], dtype=pto.f8e4m3, @@ -3757,6 +3779,8 @@ def main() -> None: "mte_gm_l1_frac", "mte_l1_bt", "mte_l1_fb", + "load_cbuf_to_ca", + "load_cbuf_to_cb", "vldsx2", "vldas", "vldus", @@ -6813,6 +6837,17 @@ def _enter_inline_simt_with_resource_attr(): expect("pto.vsstb" in vsstb_post_update_surface_text, "vsstb(..., post_update=ON) should still lower through pto.vsstb on the current VPTO IR") expect("-> !pto.ptr" in vsstb_post_update_surface_text, "vsstb(..., post_update=ON) should request the updated destination pointer result") expect("pto.mte_l1_l0b" in public_surface_text, "mte_l1_l0b(...) should lower to pto.mte_l1_l0b") + expect(public_surface_text.count("pto.load_cbuf_to_ca") == 1, "load_cbuf_to_ca(...) should lower directly to pto.load_cbuf_to_ca") + expect(public_surface_text.count("pto.load_cbuf_to_cb") == 1, "load_cbuf_to_cb(...) should lower directly to pto.load_cbuf_to_cb") + explicit_ca_controls = ( + r"pto\.load_cbuf_to_ca .*%c4_i64(?:_\d+)?, %c0_i64(?:_\d+)?, " + r"%c4_i64(?:_\d+)?, %c16_i64(?:_\d+)?, " + r"%c16_i64(?:_\d+)?, %c4_i64(?:_\d+)? \{transpose = true\}" + ) + expect( + re.search(explicit_ca_controls, public_surface_text) is not None, + "explicit L1-to-L0 controls should preserve independent source and destination strides", + ) expect("pto.mte_l1_l0a_mx" in public_surface_text, "mte_l1_l0a_mx(...) should lower to pto.mte_l1_l0a_mx") expect("pto.mte_l1_l0b_mx" in public_surface_text, "mte_l1_l0b_mx(...) should lower to pto.mte_l1_l0b_mx") expect( From fb4bc312a34bf1ae4fce6aaa47aa4aa257277c83 Mon Sep 17 00:00:00 2001 From: andodo Date: Wed, 5 Aug 2026 10:15:12 +0800 Subject: [PATCH 012/122] feat(ptodsl): unify L1 to L0 load APIs --- .../docs/user_guide/07-data-movement-ops.md | 20 +-- ptodsl/ptodsl/_ops.py | 128 +++++++++--------- ptodsl/ptodsl/pto.py | 1 - ptodsl/tests/test_jit_compile.py | 20 ++- 4 files changed, 87 insertions(+), 82 deletions(-) diff --git a/ptodsl/docs/user_guide/07-data-movement-ops.md b/ptodsl/docs/user_guide/07-data-movement-ops.md index 042d3474e5..91c55a3997 100644 --- a/ptodsl/docs/user_guide/07-data-movement-ops.md +++ b/ptodsl/docs/user_guide/07-data-movement-ops.md @@ -939,12 +939,12 @@ Cube compute step; it does not issue those transfers itself. ### Operand loading: L1 → L0A / L0B -#### `pto.load_cbuf_to_ca(src: PtrType, dst: PtrType, m_start: int, k_start: int, m_step: int, k_step: int, src_stride: int, dst_stride: int, *, transpose: bool = False) -> None` -#### `pto.load_cbuf_to_cb(src: PtrType, dst: PtrType, m_start: int, k_start: int, m_step: int, k_step: int, src_stride: int, dst_stride: int, *, transpose: bool = False) -> None` +#### `pto.mte_l1_l0a(src: PtrType, dst: PtrType, m_start: int, k_start: int, m_step: int, k_step: int, src_stride: int, dst_stride: int, *, transpose: bool = False) -> None` +#### `pto.mte_l1_l0b(src: PtrType, dst: PtrType, m_start: int, k_start: int, m_step: int, k_step: int, src_stride: int, dst_stride: int, *, transpose: bool = False) -> None` -**Description**: Explicit-control L1-to-L0A/L0B loads. Unlike the structured -`mte_l1_l0a/b` wrappers, these APIs preserve the authored fractal-block control -fields and do not infer strides from a logical tile shape. +**Description**: Explicit-control L1-to-L0A/L0B loads. This overload preserves +the authored fractal-block control fields and does not infer strides from a +logical tile shape. **Parameters**: @@ -959,9 +959,9 @@ fields and do not infer strides from a logical tile shape. **Returns**: None (side-effect operation). -Use these APIs when a source/destination subregion has layout control that cannot -be reconstructed from a canonical `m`/`k`/`n` tile shape. The pointer order is -always `src, dst`. +Use this overload when a source/destination subregion has layout control that +cannot be reconstructed from a canonical `m`/`k`/`n` tile shape. The pointer +order is always `src, dst`. --- @@ -1198,8 +1198,8 @@ pto.mte_l0c_gm( | GM → L1 | `mte_gm_l1` | gm | l1 | | GM → L1 (fractal) | `mte_gm_l1_frac` | gm | l1 | | L1 → UB | `mte_l1_ub` | l1 | ub | -| L1 → L0A (explicit control) | `load_cbuf_to_ca` | l1 | l0a | -| L1 → L0B (explicit control) | `load_cbuf_to_cb` | l1 | l0b | +| L1 → L0A (explicit control) | `mte_l1_l0a` | l1 | l0a | +| L1 → L0B (explicit control) | `mte_l1_l0b` | l1 | l0b | | L1 → L0A | `mte_l1_l0a` | l1 | l0a | | L1 → L0B | `mte_l1_l0b` | l1 | l0b | | L1 → L0A (MX) | `mte_l1_l0a_mx` | l1 | l0a | diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 35c3c152fd..7a06c045e3 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -5276,77 +5276,50 @@ def mem_bar(barrier_type): _pto.MemBarOp(kind=_membar_attr(barrier_name)) -@_explicit_mode_only("pto.load_cbuf_to_ca(...)") -def load_cbuf_to_ca( - source, - destination, - m_start, - k_start, - m_step, - k_step, - src_stride, - dst_stride, - *, - transpose=False, -): - """``pto.load_cbuf_to_ca`` – explicit-control L1-to-L0A load.""" - _pto.LoadCbufToCaOp( - unwrap_surface_value(source), - unwrap_surface_value(destination), - _coerce_i64(m_start, context="load_cbuf_to_ca m_start"), - _coerce_i64(k_start, context="load_cbuf_to_ca k_start"), - _coerce_i64(m_step, context="load_cbuf_to_ca m_step"), - _coerce_i64(k_step, context="load_cbuf_to_ca k_step"), - _coerce_i64(src_stride, context="load_cbuf_to_ca src_stride"), - _coerce_i64(dst_stride, context="load_cbuf_to_ca dst_stride"), - transpose=transpose, - ) - - -@_explicit_mode_only("pto.load_cbuf_to_cb(...)") -def load_cbuf_to_cb( - source, - destination, - m_start, - k_start, - m_step, - k_step, - src_stride, - dst_stride, - *, - transpose=False, -): - """``pto.load_cbuf_to_cb`` – explicit-control L1-to-L0B load.""" - _pto.LoadCbufToCbOp( - unwrap_surface_value(source), - unwrap_surface_value(destination), - _coerce_i64(m_start, context="load_cbuf_to_cb m_start"), - _coerce_i64(k_start, context="load_cbuf_to_cb k_start"), - _coerce_i64(m_step, context="load_cbuf_to_cb m_step"), - _coerce_i64(k_step, context="load_cbuf_to_cb k_step"), - _coerce_i64(src_stride, context="load_cbuf_to_cb src_stride"), - _coerce_i64(dst_stride, context="load_cbuf_to_cb dst_stride"), - transpose=transpose, - ) - - @_explicit_mode_only("pto.mte_l1_l0a(...)") def mte_l1_l0a( source, destination, - m, - k, + m_start, + k_start, + m_step=None, + k_step=None, + src_stride=None, + dst_stride=None, *, start_row=0, start_col=0, transpose=False, ): - """``pto.mte_l1_l0a`` – cube-side LEFT staging.""" + """``pto.mte_l1_l0a`` – structured or explicit-control L1-to-L0A load.""" + explicit_controls = (m_step, k_step, src_stride, dst_stride) + if any(value is not None for value in explicit_controls): + if any(value is None for value in explicit_controls): + raise TypeError( + "mte_l1_l0a(...) explicit control requires m_step, k_step, " + "src_stride, and dst_stride together" + ) + if start_row != 0 or start_col != 0: + raise TypeError( + "mte_l1_l0a(...) explicit control does not accept start_row or start_col" + ) + _pto.LoadCbufToCaOp( + unwrap_surface_value(source), + unwrap_surface_value(destination), + _coerce_i64(m_start, context="mte_l1_l0a m_start"), + _coerce_i64(k_start, context="mte_l1_l0a k_start"), + _coerce_i64(m_step, context="mte_l1_l0a m_step"), + _coerce_i64(k_step, context="mte_l1_l0a k_step"), + _coerce_i64(src_stride, context="mte_l1_l0a src_stride"), + _coerce_i64(dst_stride, context="mte_l1_l0a dst_stride"), + transpose=transpose, + ) + return _pto.MteL1L0aOp( unwrap_surface_value(source), unwrap_surface_value(destination), - _coerce_i64(m, context="mte_l1_l0a m"), - _coerce_i64(k, context="mte_l1_l0a k"), + _coerce_i64(m_start, context="mte_l1_l0a m"), + _coerce_i64(k_start, context="mte_l1_l0a k"), _coerce_i64(start_row, context="mte_l1_l0a start_row"), _coerce_i64(start_col, context="mte_l1_l0a start_col"), transpose=transpose, @@ -5357,19 +5330,46 @@ def mte_l1_l0a( def mte_l1_l0b( source, destination, - k, - n, + m_start, + k_start, + m_step=None, + k_step=None, + src_stride=None, + dst_stride=None, *, start_row=0, start_col=0, transpose=False, ): - """``pto.mte_l1_l0b`` – cube-side RIGHT staging.""" + """``pto.mte_l1_l0b`` – structured or explicit-control L1-to-L0B load.""" + explicit_controls = (m_step, k_step, src_stride, dst_stride) + if any(value is not None for value in explicit_controls): + if any(value is None for value in explicit_controls): + raise TypeError( + "mte_l1_l0b(...) explicit control requires m_step, k_step, " + "src_stride, and dst_stride together" + ) + if start_row != 0 or start_col != 0: + raise TypeError( + "mte_l1_l0b(...) explicit control does not accept start_row or start_col" + ) + _pto.LoadCbufToCbOp( + unwrap_surface_value(source), + unwrap_surface_value(destination), + _coerce_i64(m_start, context="mte_l1_l0b m_start"), + _coerce_i64(k_start, context="mte_l1_l0b k_start"), + _coerce_i64(m_step, context="mte_l1_l0b m_step"), + _coerce_i64(k_step, context="mte_l1_l0b k_step"), + _coerce_i64(src_stride, context="mte_l1_l0b src_stride"), + _coerce_i64(dst_stride, context="mte_l1_l0b dst_stride"), + transpose=transpose, + ) + return _pto.MteL1L0bOp( unwrap_surface_value(source), unwrap_surface_value(destination), - _coerce_i64(k, context="mte_l1_l0b k"), - _coerce_i64(n, context="mte_l1_l0b n"), + _coerce_i64(m_start, context="mte_l1_l0b k"), + _coerce_i64(k_start, context="mte_l1_l0b n"), _coerce_i64(start_row, context="mte_l1_l0b start_row"), _coerce_i64(start_col, context="mte_l1_l0b start_col"), transpose=transpose, diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index 27418a04b4..5cb8b96168 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -127,7 +127,6 @@ set_atomic_add, set_atomic_max, set_atomic_min, set_atomic_none, set_atomic_f32, set_atomic_f16, set_atomic_bf16, set_atomic_s32, set_atomic_s16, set_atomic_s8, - load_cbuf_to_ca, load_cbuf_to_cb, mte_l1_l0a, mte_l1_l0b, mte_l1_l0a_mx, mte_l1_l0b_mx, mte_l0c_l1, mte_l0c_gm, mte_l0c_ub, mad, mad_acc, mad_bias, mad_mx, mad_mx_acc, mad_mx_bias, diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index a962a135ea..d82aed2f6d 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -2391,7 +2391,7 @@ def public_surface_exports_probe( memory_space=pto.MemorySpace.RIGHT, valid_shape=[16, 16], ) - pto.load_cbuf_to_ca( + pto.mte_l1_l0a( lhs_tile.as_ptr(), lhs_l0a.as_ptr(), 4, @@ -2402,7 +2402,7 @@ def public_surface_exports_probe( 4, transpose=True, ) - pto.load_cbuf_to_cb( + pto.mte_l1_l0b( rhs_tile.as_ptr(), rhs_l0b.as_ptr(), 4, @@ -3779,8 +3779,6 @@ def main() -> None: "mte_gm_l1_frac", "mte_l1_bt", "mte_l1_fb", - "load_cbuf_to_ca", - "load_cbuf_to_cb", "vldsx2", "vldas", "vldus", @@ -3818,6 +3816,14 @@ def main() -> None: expect(isinstance(fake_empty, _FakeTensor), "pto.empty_like(...) should preserve host tensor factory type") expect(fake_empty.shape == fake_tensor.shape, "pto.empty_like(...) should preserve the logical tensor shape") expect(not hasattr(pto, "scalar"), "pto.scalar should not remain in the public pto namespace") + expect( + not hasattr(pto, "load_cbuf_to_ca"), + "pto.load_cbuf_to_ca should not be exported; use pto.mte_l1_l0a(...) explicit control", + ) + expect( + not hasattr(pto, "load_cbuf_to_cb"), + "pto.load_cbuf_to_cb should not be exported; use pto.mte_l1_l0b(...) explicit control", + ) expect(hasattr(pto, "tile"), "pto.tile should be exported from the public namespace") expect(hasattr(pto, "vmi"), "pto.vmi should be exported from the public namespace") expect(hasattr(pto.tile, "load"), "pto.tile.load should be exported from the public tile namespace") @@ -6837,8 +6843,8 @@ def _enter_inline_simt_with_resource_attr(): expect("pto.vsstb" in vsstb_post_update_surface_text, "vsstb(..., post_update=ON) should still lower through pto.vsstb on the current VPTO IR") expect("-> !pto.ptr" in vsstb_post_update_surface_text, "vsstb(..., post_update=ON) should request the updated destination pointer result") expect("pto.mte_l1_l0b" in public_surface_text, "mte_l1_l0b(...) should lower to pto.mte_l1_l0b") - expect(public_surface_text.count("pto.load_cbuf_to_ca") == 1, "load_cbuf_to_ca(...) should lower directly to pto.load_cbuf_to_ca") - expect(public_surface_text.count("pto.load_cbuf_to_cb") == 1, "load_cbuf_to_cb(...) should lower directly to pto.load_cbuf_to_cb") + expect(public_surface_text.count("pto.load_cbuf_to_ca") == 1, "explicit mte_l1_l0a(...) should lower directly to pto.load_cbuf_to_ca") + expect(public_surface_text.count("pto.load_cbuf_to_cb") == 1, "explicit mte_l1_l0b(...) should lower directly to pto.load_cbuf_to_cb") explicit_ca_controls = ( r"pto\.load_cbuf_to_ca .*%c4_i64(?:_\d+)?, %c0_i64(?:_\d+)?, " r"%c4_i64(?:_\d+)?, %c16_i64(?:_\d+)?, " @@ -6846,7 +6852,7 @@ def _enter_inline_simt_with_resource_attr(): ) expect( re.search(explicit_ca_controls, public_surface_text) is not None, - "explicit L1-to-L0 controls should preserve independent source and destination strides", + "explicit mte_l1_l0a controls should preserve independent source and destination strides", ) expect("pto.mte_l1_l0a_mx" in public_surface_text, "mte_l1_l0a_mx(...) should lower to pto.mte_l1_l0a_mx") expect("pto.mte_l1_l0b_mx" in public_surface_text, "mte_l1_l0b_mx(...) should lower to pto.mte_l1_l0b_mx") From b6e077ab5f654ff80f023a13d9b26d48a7fd1b45 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Wed, 5 Aug 2026 13:42:38 +0800 Subject: [PATCH 013/122] fix: lower tfillpad inplace mode correctly --- lib/PTO/Transforms/ExpandTileOp.cpp | 4 ++++ lib/TileOps/a5/_fillpad.py | 19 +++++++++++++-- ...pand_tile_op_tilelang_tfillpad_inplace.pto | 24 +++++++++++-------- 3 files changed, 35 insertions(+), 12 deletions(-) diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 27aeeacd95..50e8b3935a 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -522,6 +522,10 @@ static void appendOpContextAttrs( if (auto tci = dyn_cast(op)) { attrs.emplace_back("descending", tci.getDescending() ? "true" : "false"); } + if (auto tfillpad = dyn_cast(op)) { + attrs.emplace_back( + "mode", pto::stringifyTFillPadMode(tfillpad.getMode()).str()); + } if (auto tscatter = dyn_cast(op)) { if (auto maskPatternAttr = tscatter.getMaskPatternAttr()) { attrs.emplace_back( diff --git a/lib/TileOps/a5/_fillpad.py b/lib/TileOps/a5/_fillpad.py index 5e56896e47..7c19e5679f 100644 --- a/lib/TileOps/a5/_fillpad.py +++ b/lib/TileOps/a5/_fillpad.py @@ -148,8 +148,23 @@ def _fill(dst, row_start, row_stop, col_start, col_stop, scalar_tail_start=None) def _fill_inplace(dst, src_valid_rows, src_valid_cols, dst_valid_rows, dst_valid_cols): - _fill(dst, 0, src_valid_rows, src_valid_cols, dst_valid_cols) - _fill(dst, src_valid_rows, dst_valid_rows, 0, dst_valid_cols) + fill_scalar = _fill_scalar(dst) + # TileDSL v1 has no vstus/vstas equivalent for unaligned right padding. + # A masked vsts starting at src_valid_cols can overwrite valid elements. + with pto.for_(0, src_valid_rows, step=1) as row: + with pto.for_(src_valid_cols, dst_valid_cols, step=1) as col: + scalar.store(fill_scalar, dst[row, col]) + + lanes = pto.elements_per_vreg(dst.dtype) + scalar_tail_start = _scalar_tail_start(dst, lanes) + _fill( + dst, + src_valid_rows, + dst_valid_rows, + 0, + dst_valid_cols, + scalar_tail_start=scalar_tail_start, + ) def register_fillpad(): diff --git a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto index 93ddc14923..f4ba08303c 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto @@ -17,24 +17,28 @@ // pto.tfillpad (inplace) should be lowered to vector-style VPTO IR. // CHECK: func.func @TFILLPAD_INPLACE // CHECK-NOT: pto.tfillpad ins -// CHECK: pto.vecscope // CHECK: pto.castptr +// Unaligned right padding must use scalar stores because TileDSL v1 has no +// vstus/vstas equivalent and a masked vsts would overwrite valid elements. +// CHECK: pto.store % +// CHECK: pto.vecscope // CHECK-DAG: pto.vdup -// CHECK-DAG: pto.vlds // CHECK-DAG: pto.vsts module attributes {pto.kernel_kind = #pto.kernel_kind} { func.func @TFILLPAD_INPLACE() { - // 原地操作：src 和 dst 是同一个 Tile - // 有效区域 8x48，总容量 16x64 - %tile = pto.alloc_tile - : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf - // src 和 dst 相同，并显式选择原地填充模式 - pto.tfillpad ins(%tile : !pto.tile_buf) - outs(%tile : !pto.tile_buf) + outs(%dst : !pto.tile_buf) {mode = #pto.tfillpad_mode} return From ea78d8d04c66bbedc3387824ef7ee0acf679016a Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Wed, 5 Aug 2026 14:57:25 +0800 Subject: [PATCH 014/122] fix: mask tfillpad inplace vector stores --- lib/TileOps/a5/_fillpad.py | 30 +++++++++++++++---- ...pand_tile_op_tilelang_tfillpad_inplace.pto | 15 +++++----- 2 files changed, 33 insertions(+), 12 deletions(-) diff --git a/lib/TileOps/a5/_fillpad.py b/lib/TileOps/a5/_fillpad.py index 7c19e5679f..0623c591cb 100644 --- a/lib/TileOps/a5/_fillpad.py +++ b/lib/TileOps/a5/_fillpad.py @@ -149,13 +149,33 @@ def _fill(dst, row_start, row_stop, col_start, col_stop, scalar_tail_start=None) def _fill_inplace(dst, src_valid_rows, src_valid_cols, dst_valid_rows, dst_valid_cols): fill_scalar = _fill_scalar(dst) - # TileDSL v1 has no vstus/vstas equivalent for unaligned right padding. - # A masked vsts starting at src_valid_cols can overwrite valid elements. + dtype = dst.dtype + lanes = pto.elements_per_vreg(dtype) + cols = dst.shape[1] + dst_ptr = dst.as_ptr() + + # Keep each store base vector-aligned and mask out the source prefix. This + # matches vstus/vstas semantics without issuing an unaligned vector store. with pto.for_(0, src_valid_rows, step=1) as row: - with pto.for_(src_valid_cols, dst_valid_cols, step=1) as col: - scalar.store(fill_scalar, dst[row, col]) + dst_remained = dst_valid_cols + src_remained = src_valid_cols + col_loop = pto.for_(0, dst_valid_cols, step=lanes).carry( + dst_remained=dst_remained, + src_remained=src_remained, + ) + with col_loop: + col = col_loop.iv + dst_mask, dst_remained = pto.make_mask(dtype, dst_remained) + src_mask, src_remained = pto.make_mask(dtype, src_remained) + fill_mask = pto.pxor(dst_mask, src_mask, dst_mask) + vec = pto.vdup(fill_scalar, fill_mask) + addr = pto.addptr(dst_ptr, row * cols + col) + pto.vsts(vec, addr, 0, fill_mask) + col_loop.update( + dst_remained=dst_remained, + src_remained=src_remained, + ) - lanes = pto.elements_per_vreg(dst.dtype) scalar_tail_start = _scalar_tail_start(dst, lanes) _fill( dst, diff --git a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto index f4ba08303c..3b1adfb019 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto @@ -17,12 +17,13 @@ // pto.tfillpad (inplace) should be lowered to vector-style VPTO IR. // CHECK: func.func @TFILLPAD_INPLACE // CHECK-NOT: pto.tfillpad ins -// CHECK: pto.castptr -// Unaligned right padding must use scalar stores because TileDSL v1 has no -// vstus/vstas equivalent and a masked vsts would overwrite valid elements. -// CHECK: pto.store % +// CHECK: %[[MAX:.*]] = arith.constant 3.40282347E+38 : f32 +// Unaligned right padding uses an aligned store base and masks out the source +// prefix, matching vstus/vstas behavior without overwriting valid elements. // CHECK: pto.vecscope -// CHECK-DAG: pto.vdup +// CHECK: pto.castptr +// CHECK-DAG: pto.pxor +// CHECK-DAG: pto.vdup %[[MAX]], // CHECK-DAG: pto.vsts module attributes {pto.kernel_kind = #pto.kernel_kind} { @@ -32,14 +33,14 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { blayout=row_major, slayout=none_box, fractal=512, pad=3> %dst = pto.alloc_tile : !pto.tile_buf + blayout=row_major, slayout=none_box, fractal=512, pad=2> // Column 7 is deliberately unaligned. The runtime ST separately models // the PTO-ISA in-place calling convention with aliased tile addresses. pto.tfillpad ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) + blayout=row_major, slayout=none_box, fractal=512, pad=2>) {mode = #pto.tfillpad_mode} return } From 2823e2213378f7143059ad3d39a8fc22436a7500 Mon Sep 17 00:00:00 2001 From: and0d0 Date: Wed, 5 Aug 2026 17:52:46 +0800 Subject: [PATCH 015/122] fix(ptodsl): preserve legacy L1 to L0 keywords --- ptodsl/ptodsl/_ops.py | 56 ++++++++++++++------------------ ptodsl/tests/test_jit_compile.py | 8 ++--- 2 files changed, 28 insertions(+), 36 deletions(-) diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 7a06c045e3..188afe5031 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -5280,34 +5280,30 @@ def mem_bar(barrier_type): def mte_l1_l0a( source, destination, - m_start, - k_start, - m_step=None, - k_step=None, - src_stride=None, - dst_stride=None, - *, + m, + k, + *explicit_controls, start_row=0, start_col=0, transpose=False, ): """``pto.mte_l1_l0a`` – structured or explicit-control L1-to-L0A load.""" - explicit_controls = (m_step, k_step, src_stride, dst_stride) - if any(value is not None for value in explicit_controls): - if any(value is None for value in explicit_controls): + if explicit_controls: + if len(explicit_controls) != 4: raise TypeError( - "mte_l1_l0a(...) explicit control requires m_step, k_step, " - "src_stride, and dst_stride together" + "mte_l1_l0a(...) explicit control requires m_start, k_start, " + "m_step, k_step, src_stride, and dst_stride" ) if start_row != 0 or start_col != 0: raise TypeError( "mte_l1_l0a(...) explicit control does not accept start_row or start_col" ) + m_step, k_step, src_stride, dst_stride = explicit_controls _pto.LoadCbufToCaOp( unwrap_surface_value(source), unwrap_surface_value(destination), - _coerce_i64(m_start, context="mte_l1_l0a m_start"), - _coerce_i64(k_start, context="mte_l1_l0a k_start"), + _coerce_i64(m, context="mte_l1_l0a m_start"), + _coerce_i64(k, context="mte_l1_l0a k_start"), _coerce_i64(m_step, context="mte_l1_l0a m_step"), _coerce_i64(k_step, context="mte_l1_l0a k_step"), _coerce_i64(src_stride, context="mte_l1_l0a src_stride"), @@ -5318,8 +5314,8 @@ def mte_l1_l0a( _pto.MteL1L0aOp( unwrap_surface_value(source), unwrap_surface_value(destination), - _coerce_i64(m_start, context="mte_l1_l0a m"), - _coerce_i64(k_start, context="mte_l1_l0a k"), + _coerce_i64(m, context="mte_l1_l0a m"), + _coerce_i64(k, context="mte_l1_l0a k"), _coerce_i64(start_row, context="mte_l1_l0a start_row"), _coerce_i64(start_col, context="mte_l1_l0a start_col"), transpose=transpose, @@ -5330,34 +5326,30 @@ def mte_l1_l0a( def mte_l1_l0b( source, destination, - m_start, - k_start, - m_step=None, - k_step=None, - src_stride=None, - dst_stride=None, - *, + k, + n, + *explicit_controls, start_row=0, start_col=0, transpose=False, ): """``pto.mte_l1_l0b`` – structured or explicit-control L1-to-L0B load.""" - explicit_controls = (m_step, k_step, src_stride, dst_stride) - if any(value is not None for value in explicit_controls): - if any(value is None for value in explicit_controls): + if explicit_controls: + if len(explicit_controls) != 4: raise TypeError( - "mte_l1_l0b(...) explicit control requires m_step, k_step, " - "src_stride, and dst_stride together" + "mte_l1_l0b(...) explicit control requires m_start, k_start, " + "m_step, k_step, src_stride, and dst_stride" ) if start_row != 0 or start_col != 0: raise TypeError( "mte_l1_l0b(...) explicit control does not accept start_row or start_col" ) + m_step, k_step, src_stride, dst_stride = explicit_controls _pto.LoadCbufToCbOp( unwrap_surface_value(source), unwrap_surface_value(destination), - _coerce_i64(m_start, context="mte_l1_l0b m_start"), - _coerce_i64(k_start, context="mte_l1_l0b k_start"), + _coerce_i64(k, context="mte_l1_l0b m_start"), + _coerce_i64(n, context="mte_l1_l0b k_start"), _coerce_i64(m_step, context="mte_l1_l0b m_step"), _coerce_i64(k_step, context="mte_l1_l0b k_step"), _coerce_i64(src_stride, context="mte_l1_l0b src_stride"), @@ -5368,8 +5360,8 @@ def mte_l1_l0b( _pto.MteL1L0bOp( unwrap_surface_value(source), unwrap_surface_value(destination), - _coerce_i64(m_start, context="mte_l1_l0b k"), - _coerce_i64(k_start, context="mte_l1_l0b n"), + _coerce_i64(k, context="mte_l1_l0b k"), + _coerce_i64(n, context="mte_l1_l0b n"), _coerce_i64(start_row, context="mte_l1_l0b start_row"), _coerce_i64(start_col, context="mte_l1_l0b start_col"), transpose=transpose, diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index d82aed2f6d..da4862997f 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -2142,16 +2142,16 @@ def public_cube_surface_probe( pto.mte_l1_l0a( lhs_tile.as_ptr(), lhs_l0a.as_ptr(), - m, - k, + m=m, + k=k, start_row=start_row, start_col=start_col, ) pto.mte_l1_l0b( rhs_tile.as_ptr(), rhs_l0b.as_ptr(), - k, - n, + k=k, + n=n, start_row=start_col, start_col=start_row, transpose=True, From 1ffdfa63dfff75fbe0c1aadeedda5ffd3f54f4e3 Mon Sep 17 00:00:00 2001 From: and0d0 Date: Wed, 5 Aug 2026 21:32:55 +0800 Subject: [PATCH 016/122] fix(vpto): verify explicit L1 to L0 loads --- include/PTO/IR/VPTOOps.td | 2 + lib/PTO/IR/VPTO.cpp | 38 +++++++++++++++++ ...ad_cbuf_to_l0_control_verifier_invalid.pto | 21 ++++++++++ .../vpto/load_cbuf_to_l0_verifier_invalid.pto | 42 +++++++++++++++++++ 4 files changed, 103 insertions(+) create mode 100644 test/lit/vpto/load_cbuf_to_l0_control_verifier_invalid.pto create mode 100644 test/lit/vpto/load_cbuf_to_l0_verifier_invalid.pto diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index 5d139fe611..148a9b133e 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -1351,6 +1351,7 @@ def PTO_LoadCbufToCaOp : PTO_MteOp<"load_cbuf_to_ca"> { ); let results = (outs); + let hasVerifier = 1; let assemblyFormat = [{ $source `,` $destination `,` $m_start `,` $k_start `,` $m_step `,` @@ -1399,6 +1400,7 @@ def PTO_LoadCbufToCbOp : PTO_MteOp<"load_cbuf_to_cb"> { ); let results = (outs); + let hasVerifier = 1; let assemblyFormat = [{ $source `,` $destination `,` $m_start `,` $k_start `,` $m_step `,` diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index 6ca5446724..591c3cc8d5 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -7804,6 +7804,32 @@ static LogicalResult verifyMxDestinationAlignment(Operation *op, << kMxDestinationAddressUnitBytes << " bytes, got " << *address; } +template +static LogicalResult verifyExplicitCubeBridgeLoadControls(OpTy op) { + auto checkNonNegativeConst = [&](Value value, StringRef name) -> LogicalResult { + APInt intValue; + if (matchPattern(value, m_ConstantInt(&intValue)) && intValue.isNegative()) + return op.emitOpError() << name << " must be non-negative"; + return success(); + }; + auto checkPositiveConst = [&](Value value, StringRef name) -> LogicalResult { + APInt intValue; + if (matchPattern(value, m_ConstantInt(&intValue)) && + (intValue.isNegative() || intValue.isZero())) + return op.emitOpError() << name << " must be greater than zero"; + return success(); + }; + + if (failed(checkNonNegativeConst(op.getMStart(), "m_start")) || + failed(checkNonNegativeConst(op.getKStart(), "k_start")) || + failed(checkPositiveConst(op.getMStep(), "m_step")) || + failed(checkPositiveConst(op.getKStep(), "k_step")) || + failed(checkPositiveConst(op.getSrcStride(), "src_stride")) || + failed(checkPositiveConst(op.getDstStride(), "dst_stride"))) + return failure(); + return success(); +} + LogicalResult MteL0cL1Op::verify() { if (!isBufferLike(getSource().getType()) || !isBufferLike(getDestination().getType())) @@ -7867,6 +7893,18 @@ LogicalResult LoadCbufToCbMxOp::verify() { "y_start_position"); } +LogicalResult LoadCbufToCaOp::verify() { + if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::LEFT, "LEFT"))) + return failure(); + return verifyExplicitCubeBridgeLoadControls(*this); +} + +LogicalResult LoadCbufToCbOp::verify() { + if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::RIGHT, "RIGHT"))) + return failure(); + return verifyExplicitCubeBridgeLoadControls(*this); +} + void MteL1L0aOp::getEffects( SmallVectorImpl> &effects) { diff --git a/test/lit/vpto/load_cbuf_to_l0_control_verifier_invalid.pto b/test/lit/vpto/load_cbuf_to_l0_control_verifier_invalid.pto new file mode 100644 index 0000000000..2146592157 --- /dev/null +++ b/test/lit/vpto/load_cbuf_to_l0_control_verifier_invalid.pto @@ -0,0 +1,21 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @invalid_control_values() attributes {pto.kernel} { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %src = pto.castptr %c0 : i64 -> !pto.ptr + %dst = pto.castptr %c0 : i64 -> !pto.ptr + // CHECK: 'pto.load_cbuf_to_cb' op m_step must be greater than zero + pto.load_cbuf_to_cb %src, %dst, %c0, %c0, %c0, %c1, %c1, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + return + } +} diff --git a/test/lit/vpto/load_cbuf_to_l0_verifier_invalid.pto b/test/lit/vpto/load_cbuf_to_l0_verifier_invalid.pto new file mode 100644 index 0000000000..867edc39ad --- /dev/null +++ b/test/lit/vpto/load_cbuf_to_l0_verifier_invalid.pto @@ -0,0 +1,42 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @invalid_source_space() attributes {pto.kernel} { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %src = pto.castptr %c0 : i64 -> !pto.ptr + %dst = pto.castptr %c0 : i64 -> !pto.ptr + // CHECK: 'pto.load_cbuf_to_ca' op requires MAT source + pto.load_cbuf_to_ca %src, %dst, %c0, %c0, %c1, %c1, %c1, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + return + } + + func.func @invalid_destination_space() attributes {pto.kernel} { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %src = pto.castptr %c0 : i64 -> !pto.ptr + %dst = pto.castptr %c0 : i64 -> !pto.ptr + // CHECK: 'pto.load_cbuf_to_ca' op requires LEFT destination + pto.load_cbuf_to_ca %src, %dst, %c0, %c0, %c1, %c1, %c1, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + return + } + + func.func @invalid_control_values() attributes {pto.kernel} { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %cneg1 = arith.constant -1 : i64 + %src = pto.castptr %c0 : i64 -> !pto.ptr + %dst = pto.castptr %c0 : i64 -> !pto.ptr + // CHECK: 'pto.load_cbuf_to_cb' op m_start must be non-negative + pto.load_cbuf_to_cb %src, %dst, %cneg1, %c0, %c1, %c1, %c1, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + return + } +} From b24044817400e70d81abd9f41fddbb839dcfdeaf Mon Sep 17 00:00:00 2001 From: and0d0 Date: Wed, 5 Aug 2026 22:14:58 +0800 Subject: [PATCH 017/122] test(ptodsl): align L1 to L0 controls with MX API pattern --- .../docs/user_guide/07-data-movement-ops.md | 4 +- ptodsl/ptodsl/_ops.py | 82 +++++++---- ptodsl/tests/test_jit_compile.py | 24 ++-- ptodsl/tests/test_vector_cube_ops.py | 132 ++++++++++++++++++ ...ad_cbuf_to_l0_control_verifier_invalid.pto | 21 --- .../vpto/load_cbuf_to_l0_verifier_invalid.pto | 33 ++++- 6 files changed, 233 insertions(+), 63 deletions(-) delete mode 100644 test/lit/vpto/load_cbuf_to_l0_control_verifier_invalid.pto diff --git a/ptodsl/docs/user_guide/07-data-movement-ops.md b/ptodsl/docs/user_guide/07-data-movement-ops.md index 91c55a3997..00ba51b3e1 100644 --- a/ptodsl/docs/user_guide/07-data-movement-ops.md +++ b/ptodsl/docs/user_guide/07-data-movement-ops.md @@ -939,8 +939,8 @@ Cube compute step; it does not issue those transfers itself. ### Operand loading: L1 → L0A / L0B -#### `pto.mte_l1_l0a(src: PtrType, dst: PtrType, m_start: int, k_start: int, m_step: int, k_step: int, src_stride: int, dst_stride: int, *, transpose: bool = False) -> None` -#### `pto.mte_l1_l0b(src: PtrType, dst: PtrType, m_start: int, k_start: int, m_step: int, k_step: int, src_stride: int, dst_stride: int, *, transpose: bool = False) -> None` +#### `pto.mte_l1_l0a(src: PtrType, dst: PtrType, *, m_start: int, k_start: int, m_step: int, k_step: int, src_stride: int, dst_stride: int, transpose: bool = False) -> None` +#### `pto.mte_l1_l0b(src: PtrType, dst: PtrType, *, m_start: int, k_start: int, m_step: int, k_step: int, src_stride: int, dst_stride: int, transpose: bool = False) -> None` **Description**: Explicit-control L1-to-L0A/L0B loads. This overload preserves the authored fractal-block control fields and does not infer strides from a diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 188afe5031..9479563605 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -5280,30 +5280,45 @@ def mem_bar(barrier_type): def mte_l1_l0a( source, destination, - m, - k, - *explicit_controls, + m=None, + k=None, + *, start_row=0, start_col=0, + m_start=None, + k_start=None, + m_step=None, + k_step=None, + src_stride=None, + dst_stride=None, transpose=False, ): - """``pto.mte_l1_l0a`` – structured or explicit-control L1-to-L0A load.""" - if explicit_controls: - if len(explicit_controls) != 4: + """``pto.mte_l1_l0a`` – structured or explicit-control L1-to-L0A load. + + Use either the existing shape-derived ``m``/``k`` form or provide all six + explicit L1-to-L0A controls. + """ + controls = (m_start, k_start, m_step, k_step, src_stride, dst_stride) + has_explicit_controls = any(control is not None for control in controls) + if has_explicit_controls: + if m is not None or k is not None: raise TypeError( - "mte_l1_l0a(...) explicit control requires m_start, k_start, " - "m_step, k_step, src_stride, and dst_stride" + "mte_l1_l0a accepts either m/k or explicit controls, not both" ) if start_row != 0 or start_col != 0: raise TypeError( - "mte_l1_l0a(...) explicit control does not accept start_row or start_col" + "mte_l1_l0a start_row/start_col are unavailable with explicit controls" + ) + if any(control is None for control in controls): + raise TypeError( + "mte_l1_l0a explicit controls require m_start, k_start, " + "m_step, k_step, src_stride, and dst_stride" ) - m_step, k_step, src_stride, dst_stride = explicit_controls _pto.LoadCbufToCaOp( unwrap_surface_value(source), unwrap_surface_value(destination), - _coerce_i64(m, context="mte_l1_l0a m_start"), - _coerce_i64(k, context="mte_l1_l0a k_start"), + _coerce_i64(m_start, context="mte_l1_l0a m_start"), + _coerce_i64(k_start, context="mte_l1_l0a k_start"), _coerce_i64(m_step, context="mte_l1_l0a m_step"), _coerce_i64(k_step, context="mte_l1_l0a k_step"), _coerce_i64(src_stride, context="mte_l1_l0a src_stride"), @@ -5311,6 +5326,8 @@ def mte_l1_l0a( transpose=transpose, ) return + if m is None or k is None: + raise TypeError("mte_l1_l0a requires m and k without explicit controls") _pto.MteL1L0aOp( unwrap_surface_value(source), unwrap_surface_value(destination), @@ -5326,30 +5343,45 @@ def mte_l1_l0a( def mte_l1_l0b( source, destination, - k, - n, - *explicit_controls, + k=None, + n=None, + *, start_row=0, start_col=0, + m_start=None, + k_start=None, + m_step=None, + k_step=None, + src_stride=None, + dst_stride=None, transpose=False, ): - """``pto.mte_l1_l0b`` – structured or explicit-control L1-to-L0B load.""" - if explicit_controls: - if len(explicit_controls) != 4: + """``pto.mte_l1_l0b`` – structured or explicit-control L1-to-L0B load. + + Use either the existing shape-derived ``k``/``n`` form or provide all six + explicit L1-to-L0B controls. + """ + controls = (m_start, k_start, m_step, k_step, src_stride, dst_stride) + has_explicit_controls = any(control is not None for control in controls) + if has_explicit_controls: + if k is not None or n is not None: raise TypeError( - "mte_l1_l0b(...) explicit control requires m_start, k_start, " - "m_step, k_step, src_stride, and dst_stride" + "mte_l1_l0b accepts either k/n or explicit controls, not both" ) if start_row != 0 or start_col != 0: raise TypeError( - "mte_l1_l0b(...) explicit control does not accept start_row or start_col" + "mte_l1_l0b start_row/start_col are unavailable with explicit controls" + ) + if any(control is None for control in controls): + raise TypeError( + "mte_l1_l0b explicit controls require m_start, k_start, " + "m_step, k_step, src_stride, and dst_stride" ) - m_step, k_step, src_stride, dst_stride = explicit_controls _pto.LoadCbufToCbOp( unwrap_surface_value(source), unwrap_surface_value(destination), - _coerce_i64(k, context="mte_l1_l0b m_start"), - _coerce_i64(n, context="mte_l1_l0b k_start"), + _coerce_i64(m_start, context="mte_l1_l0b m_start"), + _coerce_i64(k_start, context="mte_l1_l0b k_start"), _coerce_i64(m_step, context="mte_l1_l0b m_step"), _coerce_i64(k_step, context="mte_l1_l0b k_step"), _coerce_i64(src_stride, context="mte_l1_l0b src_stride"), @@ -5357,6 +5389,8 @@ def mte_l1_l0b( transpose=transpose, ) return + if k is None or n is None: + raise TypeError("mte_l1_l0b requires k and n without explicit controls") _pto.MteL1L0bOp( unwrap_surface_value(source), unwrap_surface_value(destination), diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index da4862997f..1ae91875ba 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -2394,23 +2394,23 @@ def public_surface_exports_probe( pto.mte_l1_l0a( lhs_tile.as_ptr(), lhs_l0a.as_ptr(), - 4, - 0, - 4, - 16, - 16, - 4, + m_start=4, + k_start=0, + m_step=4, + k_step=16, + src_stride=16, + dst_stride=4, transpose=True, ) pto.mte_l1_l0b( rhs_tile.as_ptr(), rhs_l0b.as_ptr(), - 4, - 0, - 4, - 16, - 16, - 4, + m_start=4, + k_start=0, + m_step=4, + k_step=16, + src_stride=16, + dst_stride=4, transpose=True, ) lhs_l0a_mx = pto.alloc_tile( diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index d63699c96a..33e485ef46 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -504,6 +504,138 @@ def test_cube_variant_wrappers_dispatch_to_generated_ops(self): getattr(_ops, func_name)(*args) self.assertEqual(op_ctor.call_args.args, expected_call) + def test_mte_l1_l0_explicit_controls_dispatch_to_load_cbuf_ops(self): + source = object() + destination = object() + controls = { + "m_start": 3, + "k_start": 5, + "m_step": 16, + "k_step": 2, + "src_stride": 8, + "dst_stride": 2, + } + + with patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "_coerce_i64", side_effect=lambda value, *, context: f"{context}:{value}"): + ca_op = MagicMock() + with patch.object(_ops._pto, "LoadCbufToCaOp", ca_op): + _ops.mte_l1_l0a(source, destination, **controls, transpose=True) + self.assertEqual( + ca_op.call_args.args, + ( + source, + destination, + "mte_l1_l0a m_start:3", + "mte_l1_l0a k_start:5", + "mte_l1_l0a m_step:16", + "mte_l1_l0a k_step:2", + "mte_l1_l0a src_stride:8", + "mte_l1_l0a dst_stride:2", + ), + ) + self.assertEqual(ca_op.call_args.kwargs, {"transpose": True}) + + cb_op = MagicMock() + with patch.object(_ops._pto, "LoadCbufToCbOp", cb_op): + _ops.mte_l1_l0b(source, destination, **controls, transpose=True) + self.assertEqual( + cb_op.call_args.args, + ( + source, + destination, + "mte_l1_l0b m_start:3", + "mte_l1_l0b k_start:5", + "mte_l1_l0b m_step:16", + "mte_l1_l0b k_step:2", + "mte_l1_l0b src_stride:8", + "mte_l1_l0b dst_stride:2", + ), + ) + self.assertEqual(cb_op.call_args.kwargs, {"transpose": True}) + + def test_mte_l1_l0_legacy_forms_preserve_keyword_compatibility(self): + source = object() + destination = object() + + with patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "_coerce_i64", side_effect=lambda value, *, context: f"{context}:{value}"): + ca_op = MagicMock() + with patch.object(_ops._pto, "MteL1L0aOp", ca_op): + _ops.mte_l1_l0a( + source, + destination, + m=128, + k=64, + start_row=2, + start_col=4, + transpose=True, + ) + self.assertEqual( + ca_op.call_args.args, + ( + source, + destination, + "mte_l1_l0a m:128", + "mte_l1_l0a k:64", + "mte_l1_l0a start_row:2", + "mte_l1_l0a start_col:4", + ), + ) + self.assertEqual(ca_op.call_args.kwargs, {"transpose": True}) + + cb_op = MagicMock() + with patch.object(_ops._pto, "MteL1L0bOp", cb_op): + _ops.mte_l1_l0b( + source, + destination, + k=64, + n=128, + start_row=4, + start_col=2, + transpose=True, + ) + self.assertEqual( + cb_op.call_args.args, + ( + source, + destination, + "mte_l1_l0b k:64", + "mte_l1_l0b n:128", + "mte_l1_l0b start_row:4", + "mte_l1_l0b start_col:2", + ), + ) + self.assertEqual(cb_op.call_args.kwargs, {"transpose": True}) + + def test_mte_l1_l0_explicit_controls_reject_ambiguous_forms(self): + source = object() + destination = object() + complete_controls = { + "m_start": 3, + "k_start": 5, + "m_step": 16, + "k_step": 2, + "src_stride": 8, + "dst_stride": 2, + } + + invalid_cases = [ + lambda: _ops.mte_l1_l0a(source, destination, m_start=3), + lambda: _ops.mte_l1_l0b(source, destination, m_start=3), + lambda: _ops.mte_l1_l0a(source, destination, m=128, k=64, **complete_controls), + lambda: _ops.mte_l1_l0b(source, destination, k=64, n=128, **complete_controls), + lambda: _ops.mte_l1_l0a(source, destination, start_row=1, **complete_controls), + lambda: _ops.mte_l1_l0b(source, destination, start_col=1, **complete_controls), + ] + for invalid_call in invalid_cases: + with self.subTest(call=invalid_call): + with self.assertRaises(TypeError): + invalid_call() + + self.assertFalse(hasattr(pto, "load_cbuf_to_ca")) + self.assertFalse(hasattr(pto, "load_cbuf_to_cb")) + def test_mad_option_wrappers_dispatch_to_generated_ops(self): lhs = object() rhs = object() diff --git a/test/lit/vpto/load_cbuf_to_l0_control_verifier_invalid.pto b/test/lit/vpto/load_cbuf_to_l0_control_verifier_invalid.pto deleted file mode 100644 index 2146592157..0000000000 --- a/test/lit/vpto/load_cbuf_to_l0_control_verifier_invalid.pto +++ /dev/null @@ -1,21 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s - -module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { - func.func @invalid_control_values() attributes {pto.kernel} { - %c0 = arith.constant 0 : i64 - %c1 = arith.constant 1 : i64 - %src = pto.castptr %c0 : i64 -> !pto.ptr - %dst = pto.castptr %c0 : i64 -> !pto.ptr - // CHECK: 'pto.load_cbuf_to_cb' op m_step must be greater than zero - pto.load_cbuf_to_cb %src, %dst, %c0, %c0, %c0, %c1, %c1, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 - return - } -} diff --git a/test/lit/vpto/load_cbuf_to_l0_verifier_invalid.pto b/test/lit/vpto/load_cbuf_to_l0_verifier_invalid.pto index 867edc39ad..2145833866 100644 --- a/test/lit/vpto/load_cbuf_to_l0_verifier_invalid.pto +++ b/test/lit/vpto/load_cbuf_to_l0_verifier_invalid.pto @@ -6,37 +6,62 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s +// RUN: split-file %s %t +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_source.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-SOURCE +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_destination.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-DESTINATION +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/cb_negative_m_start.pto -o - 2>&1 | FileCheck %s --check-prefix=CB-NEGATIVE-M-START +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/cb_zero_m_step.pto -o - 2>&1 | FileCheck %s --check-prefix=CB-ZERO-M-STEP +//--- ca_source.pto module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { func.func @invalid_source_space() attributes {pto.kernel} { %c0 = arith.constant 0 : i64 %c1 = arith.constant 1 : i64 %src = pto.castptr %c0 : i64 -> !pto.ptr %dst = pto.castptr %c0 : i64 -> !pto.ptr - // CHECK: 'pto.load_cbuf_to_ca' op requires MAT source pto.load_cbuf_to_ca %src, %dst, %c0, %c0, %c1, %c1, %c1, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 return } +} +//--- ca_destination.pto +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { func.func @invalid_destination_space() attributes {pto.kernel} { %c0 = arith.constant 0 : i64 %c1 = arith.constant 1 : i64 %src = pto.castptr %c0 : i64 -> !pto.ptr %dst = pto.castptr %c0 : i64 -> !pto.ptr - // CHECK: 'pto.load_cbuf_to_ca' op requires LEFT destination pto.load_cbuf_to_ca %src, %dst, %c0, %c0, %c1, %c1, %c1, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 return } +} +//--- cb_negative_m_start.pto +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { func.func @invalid_control_values() attributes {pto.kernel} { %c0 = arith.constant 0 : i64 %c1 = arith.constant 1 : i64 %cneg1 = arith.constant -1 : i64 %src = pto.castptr %c0 : i64 -> !pto.ptr %dst = pto.castptr %c0 : i64 -> !pto.ptr - // CHECK: 'pto.load_cbuf_to_cb' op m_start must be non-negative pto.load_cbuf_to_cb %src, %dst, %cneg1, %c0, %c1, %c1, %c1, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 return } } + +//--- cb_zero_m_step.pto +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @invalid_control_values() attributes {pto.kernel} { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %src = pto.castptr %c0 : i64 -> !pto.ptr + %dst = pto.castptr %c0 : i64 -> !pto.ptr + pto.load_cbuf_to_cb %src, %dst, %c0, %c0, %c0, %c1, %c1, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + return + } +} + +// CA-SOURCE: 'pto.load_cbuf_to_ca' op requires MAT source +// CA-DESTINATION: 'pto.load_cbuf_to_ca' op requires LEFT destination +// CB-NEGATIVE-M-START: 'pto.load_cbuf_to_cb' op m_start must be non-negative +// CB-ZERO-M-STEP: 'pto.load_cbuf_to_cb' op m_step must be greater than zero From 151e76cbcca4367f1a6380bad73fe9bfce2bb991 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Thu, 6 Aug 2026 10:43:10 +0800 Subject: [PATCH 018/122] fix: infer tfillpad lowering after memory planning --- docs/PTO_IR_manual.md | 17 ++- ...est-first-fit-four-gates-memplan-design.md | 4 +- ...ptoas-tile-native-mainline-op-migration.md | 2 +- docs/isa/tile-op/12-fill-and-padding-ops.md | 28 +++-- .../release/PTO-tile-Instruction-SPEC-v0.4.md | 25 +++-- include/PTO/IR/PTOAttrs.td | 13 --- include/PTO/IR/PTOOps.td | 5 +- include/pto-c/Dialect/PTO.h | 3 - lib/Bindings/Python/PTOModule.cpp | 19 ---- lib/CAPI/Dialect/PTO.cpp | 15 --- lib/PTO/IR/PTO.cpp | 43 ++++---- lib/PTO/Transforms/ExpandTileOp.cpp | 53 ++++++--- lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 7 +- lib/PTO/Transforms/PTOToEmitC.cpp | 25 +++-- lib/PTO/Transforms/Utils.cpp | 101 +++++++++++++++++- lib/PTO/Transforms/Utils.h | 10 ++ lib/TileOps/a5/_fillpad.py | 8 +- ptodsl/ptodsl/_ops.py | 24 +---- ptodsl/tests/test_vector_cube_ops.py | 15 ++- python/pto/dialects/pto.py | 4 - test/lit/pto/fillpad_tile_native.pto | 12 ++- .../lit/pto/movement_metadata_tile_native.pto | 4 +- .../pto/tfillpad_inplace_alias_lowering.pto | 24 ++++- .../pto/tfillpad_non_normal_mat_invalid.pto | 6 +- .../pto/tfillpad_plan_memory_inference.pto | 22 ++++ .../tfillpad_same_ssa_lowers_to_tfillpad.pto | 3 +- ...xpand_tile_op_tilelang_tfillpad_expand.pto | 5 +- ...pand_tile_op_tilelang_tfillpad_inplace.pto | 12 +-- test/samples/Fillpad/fillpad_expand.py | 4 +- .../samples/Fillpad/fillpad_expand_invalid.py | 4 +- .../fillpad_expand_pad_null_invalid.py | 4 +- test/samples/Fillpad/fillpad_inplace.py | 4 +- test/samples/runop.sh | 9 +- .../tfillpad_expand/tfillpad_expand.pto | 16 ++- .../tfillpad_inplace/tfillpad_inplace.pto | 4 +- .../tfillpad_expand/tfillpad_expand.pto | 20 ++-- .../tfillpad_inplace/tfillpad_inplace.pto | 3 +- 37 files changed, 326 insertions(+), 251 deletions(-) create mode 100644 test/lit/pto/tfillpad_plan_memory_inference.pto diff --git a/docs/PTO_IR_manual.md b/docs/PTO_IR_manual.md index 7a953c5a13..186507ee1a 100644 --- a/docs/PTO_IR_manual.md +++ b/docs/PTO_IR_manual.md @@ -8053,7 +8053,7 @@ pto.textract ins(%src, %row, %col : !pto.tile_buf<...>, index, index fp %fp : !p ##### `pto.tfillpad` - Fill Padding Region -**Summary:** Unified normal, in-place, and expand padding operation. `mode` defaults to `normal` and is never inferred from SSA aliasing or shapes. +**Summary:** Unified padding operation. PTOAS infers normal, in-place, or expand lowering from physical shapes and post-PlanMemory addresses. **Semantics:** @@ -8069,7 +8069,6 @@ expand: copy src into a possibly larger dst, then fill the expanded region |------|------|-------------| | `src` | `pto.tile_buf` | Source tile | | `dst` | `pto.tile_buf` | Destination tile (with pad config) | -| `mode` | `#pto.tfillpad_mode` | PTO-ISA execution mode; defaults to `normal` | | `padValue` | `#pto.pad_value<...>` (optional) | Explicit `TFILLPAD` template argument for `loc=mat`. When present, it must match `dst`'s tile pad configuration. | **Results:** None. Writes into `dst` via DPS pattern. @@ -8078,17 +8077,19 @@ expand: copy src into a possibly larger dst, then fill the expanded region - `dst.pad` must not be `null`. - `src` and `dst` element sizes must match, and the element size must be `1`, `2`, or `4` bytes. -- Normal and in-place modes require equal source and destination static shapes. -- Expand mode requires `dst.rows >= src.rows` and `dst.cols >= src.cols`. -- Non-normal modes require both operands to use `loc=vec`. -- If `padValue` is present, mode must be normal, `dst` must be `loc=mat`, and `padValue` must equal the tile type's `pad`. +- If source and destination physical shapes differ, every destination dimension must be at least the corresponding source dimension; PTOAS then infers expand lowering. +- If physical shapes are equal, exact starting-address equality after PlanMemory selects in-place lowering; otherwise PTOAS selects normal lowering. +- `valid_shape` does not participate in expand inference. +- In-place and expand lowering require both operands to use `loc=vec`. +- If `padValue` is present, `dst` must be `loc=mat`, and `padValue` must equal the tile type's `pad`. +- MAT always uses Normal lowering, including when source and destination share the same starting address. - For `loc=mat`, `src` and `dst` must be lowerable to the same `TFILLPAD` tile specialization, i.e. `validShape` and `pad` must be identical. **Hardware Mapping:** - VEC forms execute on the **Vector pipeline** (`PIPE_V`). - The normal homogeneous MAT form executes on `PIPE_MTE1`. -- Normal lowers to `TFILLPAD(dst, src)`; non-normal modes lower one-to-one to `TFILLPAD(dst, src)`. +- Normal lowers to `TFILLPAD(dst, src)`; compiler-inferred in-place and expand forms lower to `TFILLPAD(dst, src)`. **Basic Example:** @@ -8097,11 +8098,9 @@ pto.tfillpad ins(%src : !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) pto.tfillpad ins(%tile : !pto.tile_buf) outs(%tile : !pto.tile_buf) - {mode = #pto.tfillpad_mode} pto.tfillpad ins(%src_small : !pto.tile_buf) outs(%dst_large : !pto.tile_buf) - {mode = #pto.tfillpad_mode} ``` --- diff --git a/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md b/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md index 4750cae6fb..c4500220b7 100644 --- a/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md +++ b/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md @@ -452,7 +452,7 @@ if opPolicy.notInplaceSafe: pto.ttrans pto.tgather pto.tands / pto.tors / pto.txors -pto.tfillpad {mode = #pto.tfillpad_mode} +pto.tfillpad // dst physical shape is larger than src pto.tfmod / pto.tfmods pto.trecip / pto.trsqrt pto.trowmax / pto.trowmin / pto.trowsum / pto.trowprod @@ -461,7 +461,7 @@ pto.tcolargmax / pto.tcolargmin pto.tsort32 / pto.tmrgsort ``` -其中 `pto.tands` / `pto.tors` / `pto.txors` 和 expand 模式的 `pto.tfillpad` 是 PTOAS 侧额外保守标记的 non-inplace-safe op。它们虽然不是 scratch-output conflict，但后端/ISA 语义没有明确承诺 input/output alias 安全，memplan 不应通过地址复用隐式把它们变成 inplace 执行。 +其中 `pto.tands` / `pto.tors` / `pto.txors` 和推导为 expand lowering 的 `pto.tfillpad` 是 PTOAS 侧额外保守标记的 non-inplace-safe op。它们虽然不是 scratch-output conflict，但后端/ISA 语义没有明确承诺 input/output alias 安全，memplan 不应通过地址复用隐式把它们变成 inplace 执行。 **适用场景 sample：算法本身不支持 input/output alias。** diff --git a/docs/designs/ptoas-tile-native-mainline-op-migration.md b/docs/designs/ptoas-tile-native-mainline-op-migration.md index bfa5e3050c..3ff77473e9 100644 --- a/docs/designs/ptoas-tile-native-mainline-op-migration.md +++ b/docs/designs/ptoas-tile-native-mainline-op-migration.md @@ -187,7 +187,7 @@ PTO tile/view IR | `pto.texpands` | shape扩展和 scalar operand | | `pto.textract` | tile role、offset、fp tile 地址空间和可选 pre-quant | | `pto.tinsert` | materialize pass 当前对 tile config 有特殊推断；包含 fp/pre-quant tile role，目标是由 result type 完整携带 | -| `pto.tfillpad` | 显式 mode、src/dst alias、MemoryEffects 和 A5 MAT/PIPE 选择 | +| `pto.tfillpad` | 基于 physical shape 和 PlanMemory 地址推导 lowering、src/dst alias、MemoryEffects 和 A5 MAT/PIPE 选择 | | `pto.tsetval` | tile writer 和 result type | | `pto.tgetval` | tile reader和 scalar result | | `pto.tgather` | optional tmp、compare/index form 和 sync macro model | diff --git a/docs/isa/tile-op/12-fill-and-padding-ops.md b/docs/isa/tile-op/12-fill-and-padding-ops.md index 0900d2e179..f38c6fe02b 100644 --- a/docs/isa/tile-op/12-fill-and-padding-ops.md +++ b/docs/isa/tile-op/12-fill-and-padding-ops.md @@ -15,9 +15,8 @@ The destination tile's `pad` / `pad_value` configuration determines which value ```mlir pto.tfillpad ins(%src : !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) - {mode = #pto.tfillpad_mode} ``` -- **semantics:** the `mode` attribute selects normal, in-place, or expand behavior. It defaults to `normal`; PTOAS does not infer it from aliasing or shape. +- **semantics:** PTOAS infers normal, in-place, or expand behavior from the physical tile shapes and the addresses produced by memory planning. Users do not specify a mode. **Parameter Table:** @@ -25,24 +24,25 @@ pto.tfillpad ins(%src : !pto.tile_buf<...>) |-----------|------|-------------| | `src` | `pto.tile_buf` | Source tile. | | `dst` | `pto.tile_buf` | Destination tile carrying the pad configuration. | -| `mode` | `#pto.tfillpad_mode` | ISA mode; defaults to `normal`. | -| `padValue` | `#pto.pad_value<...>` (optional) | Explicit MAT `TFILLPAD` argument; only valid in normal mode. | +| `padValue` | `#pto.pad_value<...>` (optional) | Explicit MAT `TFILLPAD` argument. | -**Mode Table:** +**Inference Table:** -| Mode | Behavior | PTO-ISA mapping | -|------|----------|-----------------| -| `normal` | Copy valid data from `src`, then fill padding in `dst`. | `TFILLPAD(dst, src)` | -| `in_place` | Skip the copy phase and fill padding on shared storage. | `TFILLPAD(dst, src)` | -| `expand` | Copy `src` into a destination whose static shape may be larger, then fill the expanded region. | `TFILLPAD(dst, src)` | +| Compiler condition | Behavior | PTO-ISA mapping | +|--------------------|----------|-----------------| +| VEC, equal physical shapes, and different or unprovable addresses | Copy valid data from `src`, then fill padding in `dst`. | `TFILLPAD(dst, src)` | +| VEC, equal physical shapes, and identical starting addresses after memory planning | Skip the copy phase and fill padding on shared storage. | `TFILLPAD(dst, src)` | +| VEC, every `dst` physical dimension is at least the corresponding `src` dimension, and at least one is larger | Copy `src` into the larger destination and fill the expanded region. | `TFILLPAD(dst, src)` | +| Supported non-VEC form, regardless of address equality | Use the architecture's normal overload. | `TFILLPAD(dst, src)` | **Constraints:** - Source and destination element types must be compatible. - The destination tile must carry a meaningful pad configuration. -- `in_place` and `expand` are VEC-only. Normal mode also supports the homogeneous MAT overload. -- Normal and in-place modes require equal source and destination static shapes. -- Expand mode requires each destination static dimension to be greater than or equal to the source dimension. +- In-place and expand lowering are VEC-only. Normal lowering also supports the homogeneous MAT overload. +- Expand inference compares physical `shape`, not `valid_shape`. +- When physical shapes are equal, PTOAS compares exact starting addresses after PlanMemory. If equality cannot be proven, it conservatively chooses normal lowering. +- MAT always uses Normal lowering, including when source and destination share the same starting address. **Example:** @@ -52,9 +52,7 @@ pto.tfillpad ins(%src : !pto.tile_buf) pto.tfillpad ins(%tile : !pto.tile_buf) outs(%tile : !pto.tile_buf) - {mode = #pto.tfillpad_mode} pto.tfillpad ins(%src_small : !pto.tile_buf) outs(%dst_large : !pto.tile_buf) - {mode = #pto.tfillpad_mode} ``` diff --git a/docs/release/PTO-tile-Instruction-SPEC-v0.4.md b/docs/release/PTO-tile-Instruction-SPEC-v0.4.md index a3448620dc..f7e591c6e6 100644 --- a/docs/release/PTO-tile-Instruction-SPEC-v0.4.md +++ b/docs/release/PTO-tile-Instruction-SPEC-v0.4.md @@ -1769,9 +1769,8 @@ The destination tile's `pad` / `pad_value` configuration determines which value ```mlir pto.tfillpad ins(%src : !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) - {mode = #pto.tfillpad_mode} ``` -- **semantics:** `mode` selects normal, in-place, or expand behavior and defaults to `normal`. PTOAS does not infer the mode from shapes or aliasing. +- **semantics:** PTOAS infers normal, in-place, or expand behavior from physical tile shapes and post-PlanMemory addresses. Users do not specify a mode. **Parameter Table:** @@ -1779,22 +1778,24 @@ pto.tfillpad ins(%src : !pto.tile_buf<...>) |-----------|------|-------------| | `src` | `pto.tile_buf` | Source tile. | | `dst` | `pto.tile_buf` | Destination tile carrying the pad configuration. | -| `mode` | `#pto.tfillpad_mode` | PTO-ISA execution mode. | -**Mode Table:** +**Inference Table:** -| Mode | Behavior | PTO-ISA mapping | -|------|----------|-----------------| -| `normal` | Copy valid data, then fill padding. | `TFILLPAD(dst, src)` | -| `in_place` | Skip the copy phase and fill padding on shared storage. | `TFILLPAD(dst, src)` | -| `expand` | Copy into a possibly larger destination and fill the expanded region. | `TFILLPAD(dst, src)` | +| Compiler condition | Behavior | PTO-ISA mapping | +|--------------------|----------|-----------------| +| VEC, equal physical shapes, and different or unprovable addresses | Copy valid data, then fill padding. | `TFILLPAD(dst, src)` | +| VEC, equal physical shapes, and identical starting addresses after memory planning | Skip the copy phase and fill padding on shared storage. | `TFILLPAD(dst, src)` | +| VEC, destination physical shape is at least the source shape in every dimension and larger in at least one | Copy into the larger destination and fill the expanded region. | `TFILLPAD(dst, src)` | +| Supported non-VEC form, regardless of address equality | Use the architecture's normal overload. | `TFILLPAD(dst, src)` | **Constraints:** - Source and destination element types must be compatible. - The destination tile must carry a meaningful pad configuration. -- Non-normal modes are VEC-only. Normal mode also supports the homogeneous MAT overload. -- Normal and in-place modes require equal static shapes; expand requires each destination dimension to be greater than or equal to the source dimension. +- In-place and expand lowering are VEC-only. Normal lowering also supports the homogeneous MAT overload. +- Expand inference compares physical `shape`, not `valid_shape`. +- Equal physical shapes use exact starting-address equality after PlanMemory to select in-place lowering; an unprovable address relationship selects normal lowering. +- MAT always uses Normal lowering, including when source and destination share the same starting address. **Example:** @@ -1804,9 +1805,7 @@ pto.tfillpad ins(%src : !pto.tile_buf) pto.tfillpad ins(%tile : !pto.tile_buf) outs(%tile : !pto.tile_buf) - {mode = #pto.tfillpad_mode} pto.tfillpad ins(%src_small : !pto.tile_buf) outs(%dst_large : !pto.tile_buf) - {mode = #pto.tfillpad_mode} ``` diff --git a/include/PTO/IR/PTOAttrs.td b/include/PTO/IR/PTOAttrs.td index fdf95e0cea..d597173105 100644 --- a/include/PTO/IR/PTOAttrs.td +++ b/include/PTO/IR/PTOAttrs.td @@ -1190,19 +1190,6 @@ def PTO_PadValueAttr : PTO_Attr<"PadValue", "pad_value"> { let assemblyFormat = "`<` params `>`"; } -def PTO_TFillPadModeEnum : PTO_I32Enum< - "TFillPadMode", "PTO TFILLPAD execution mode", [ - I32EnumAttrCase<"Normal", 0, "normal">, - I32EnumAttrCase<"InPlace", 1, "in_place">, - I32EnumAttrCase<"Expand", 2, "expand"> - ]>; - -def PTO_TFillPadModeAttr - : EnumAttr { - let assemblyFormat = "`<` params `>`"; - let summary = "TFILLPAD normal, in-place, or expand execution mode"; -} - def PTO_CompactMode_Enum : PTO_I32Enum<"CompactMode", "Tile compact mode", [ I32EnumAttrCase<"Null", 0, "null">, I32EnumAttrCase<"Normal", 1, "normal">, diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index caad7be0d9..3a43a9804c 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -4697,13 +4697,12 @@ def TFillPadOp : PTO_TOp<"tfillpad", [ OpPipeInterface, DeclareOpInterfaceMethods ]> { - let summary = "Fill padding in normal, in-place, or expand mode (tilebuf, DPS)"; + let summary = "Fill padding with compiler-inferred lowering semantics (tilebuf, DPS)"; let arguments = (ins PTODpsType:$src, PTODpsType:$dst, - OptionalAttr:$padValue, - DefaultValuedAttr:$mode + OptionalAttr:$padValue ); let results = (outs); diff --git a/include/pto-c/Dialect/PTO.h b/include/pto-c/Dialect/PTO.h index e81917467f..cb653c2384 100644 --- a/include/pto-c/Dialect/PTO.h +++ b/include/pto-c/Dialect/PTO.h @@ -109,9 +109,6 @@ MLIR_CAPI_EXPORTED int32_t mlirPTOPadValueAttrGetValue(MlirAttribute attr); MLIR_CAPI_EXPORTED bool mlirPTOAttrIsACompactModeAttr(MlirAttribute attr); MLIR_CAPI_EXPORTED MlirAttribute mlirPTOCompactModeAttrGet(MlirContext ctx, int32_t value); MLIR_CAPI_EXPORTED int32_t mlirPTOCompactModeAttrGetValue(MlirAttribute attr); -MLIR_CAPI_EXPORTED bool mlirPTOAttrIsATFillPadModeAttr(MlirAttribute attr); -MLIR_CAPI_EXPORTED MlirAttribute mlirPTOTFillPadModeAttrGet(MlirContext ctx, int32_t value); -MLIR_CAPI_EXPORTED int32_t mlirPTOTFillPadModeAttrGetValue(MlirAttribute attr); MLIR_CAPI_EXPORTED bool mlirPTOAttrIsAAccToVecModeAttr(MlirAttribute attr); MLIR_CAPI_EXPORTED MlirAttribute mlirPTOAccToVecModeAttrGet(MlirContext ctx, int32_t value); MLIR_CAPI_EXPORTED int32_t mlirPTOAccToVecModeAttrGetValue(MlirAttribute attr); diff --git a/lib/Bindings/Python/PTOModule.cpp b/lib/Bindings/Python/PTOModule.cpp index be8576fdcc..94d3f728f4 100644 --- a/lib/Bindings/Python/PTOModule.cpp +++ b/lib/Bindings/Python/PTOModule.cpp @@ -257,12 +257,6 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { .value("DualModeSplitN", mlir::pto::AccToVecMode::DualModeSplitN) .export_values(); - py::enum_(m, "TFillPadMode") - .value("Normal", mlir::pto::TFillPadMode::Normal) - .value("InPlace", mlir::pto::TFillPadMode::InPlace) - .value("Expand", mlir::pto::TFillPadMode::Expand) - .export_values(); - py::enum_(m, "TInsertMode") .value("SPLIT2", mlir::pto::TInsertMode::SPLIT2) .value("SPLIT4", mlir::pto::TInsertMode::SPLIT4) @@ -399,19 +393,6 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { }, py::arg("cls"), py::arg("value"), py::arg("context") = py::none()); - mlir_attribute_subclass(m, "TFillPadModeAttr", - [](MlirAttribute a) -> bool { - return mlirPTOAttrIsATFillPadModeAttr(a); - }) - .def_classmethod( - "get", - [](py::object cls, mlir::pto::TFillPadMode value, MlirContext ctx) -> py::object { - MlirAttribute a = mlirPTOTFillPadModeAttrGet(ctx, static_cast(value)); - if (mlirAttributeIsNull(a)) return py::none(); - return cls(a); - }, - py::arg("cls"), py::arg("value"), py::arg("context") = py::none()); - mlir_attribute_subclass(m, "TInsertModeAttr", [](MlirAttribute a) -> bool { return mlirPTOAttrIsATInsertModeAttr(a); diff --git a/lib/CAPI/Dialect/PTO.cpp b/lib/CAPI/Dialect/PTO.cpp index 9bea9f9433..40520b90fc 100644 --- a/lib/CAPI/Dialect/PTO.cpp +++ b/lib/CAPI/Dialect/PTO.cpp @@ -734,21 +734,6 @@ int32_t mlirPTOCompactModeAttrGetValue(MlirAttribute attr) { return static_cast(a.getValue()); } -bool mlirPTOAttrIsATFillPadModeAttr(MlirAttribute attr) { - return mlir::isa(unwrap(attr)); -} - -MlirAttribute mlirPTOTFillPadModeAttrGet(MlirContext ctx, int32_t value) { - auto *c = unwrap(ctx); - return wrap(mlir::pto::TFillPadModeAttr::get( - c, static_cast(value))); -} - -int32_t mlirPTOTFillPadModeAttrGetValue(MlirAttribute attr) { - auto a = mlir::cast(unwrap(attr)); - return static_cast(a.getValue()); -} - bool mlirPTOAttrIsAAccToVecModeAttr(MlirAttribute attr) { return mlir::isa(unwrap(attr)); } diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 30d162ce5e..fb4baa1909 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -7113,8 +7113,7 @@ static bool isA5VectorPreQuantTypePair(Type srcElem, Type dstElem) { } static mlir::LogicalResult verifyTFillPadLike(Operation *op, Type srcTy, - Type dstTy, - pto::TFillPadMode mode) { + Type dstTy) { if (!isPTOShapedLike(srcTy) || !isPTOShapedLike(dstTy)) return op->emitError("expects src/dst to be PTO shaped-like types"); @@ -7142,11 +7141,23 @@ static mlir::LogicalResult verifyTFillPadLike(Operation *op, Type srcTy, auto srcSpace = getPTOMemorySpaceEnum(srcTy); auto dstSpace = getPTOMemorySpaceEnum(dstTy); - if (mode != pto::TFillPadMode::Normal && + + bool expanded = false; + for (auto [srcDim, dstDim] : llvm::zip_equal(srcShape, dstShape)) { + if (srcDim == dstDim) + continue; + if (ShapedType::isDynamic(srcDim) || ShapedType::isDynamic(dstDim)) + return op->emitError("cannot infer TFILLPAD lowering from mismatched " + "dynamic physical shapes"); + if (srcDim > dstDim) + return op->emitError( + "expects each dst physical shape dimension to be >= src"); + expanded = true; + } + if (expanded && (!srcSpace || !dstSpace || *srcSpace != pto::AddressSpace::VEC || *dstSpace != pto::AddressSpace::VEC)) - return op->emitError() - << "expects non-normal TFILLPAD mode only for loc=vec"; + return op->emitError("expects expanded TFILLPAD only for loc=vec"); // pto.tfillpad lowers to TFILLPAD(dst, src). For loc=mat, pto-isa only // exposes the homogeneous overload, so src/dst must use the same Tile<...> @@ -7192,29 +7203,19 @@ static mlir::LogicalResult verifyTFillPadLike(Operation *op, Type srcTy, return op->emitError("expects dst PadVal != Null for tfillpad"); } - if (mode != pto::TFillPadMode::Expand) { - if (srcShape != dstShape) - return op->emitError("expects src and dst to have the same static shape " - "unless mode is expand"); - return mlir::success(); - } - - if (srcShape[0] > dstShape[0] || srcShape[1] > dstShape[1]) { - return op->emitError( - "expects dst static shape to be >= src static shape for expand mode"); - } - return mlir::success(); } mlir::LogicalResult mlir::pto::TFillPadOp::verify() { - if (failed(verifyTFillPadLike(getOperation(), getSrc().getType(), getDst().getType(), - getMode()))) + if (getOperation()->getAttr("mode")) + return emitOpError("does not accept 'mode'; PTOAS infers TFILLPAD lowering " + "from physical shape and planned addresses"); + + if (failed(verifyTFillPadLike(getOperation(), getSrc().getType(), + getDst().getType()))) return failure(); if (auto padValueAttr = getPadValueAttr()) { - if (getMode() != pto::TFillPadMode::Normal) - return emitOpError("expects padValue attribute only for normal mode"); auto dstSpace = getPTOMemorySpaceEnum(getDst().getType()); if (!dstSpace || *dstSpace != pto::AddressSpace::MAT) return emitOpError("expects padValue attribute only for loc=mat tfillpad"); diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 6b9b1d3500..1643e529f3 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -32,6 +32,7 @@ #include "PTO/Support/PythonExecutable.h" #include "PTO/Transforms/Passes.h" #include "PTO/Transforms/TileOpExpansionUtils.h" +#include "Utils.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" @@ -482,7 +483,7 @@ static std::string getTRandomRoundsString(pto::TRandomOp op) { return std::to_string(op.getRounds()); } -static void appendOpContextAttrs( +static LogicalResult appendOpContextAttrs( Operation *op, SmallVectorImpl> &attrs) { if (auto tcvt = dyn_cast(op)) { @@ -527,8 +528,25 @@ static void appendOpContextAttrs( attrs.emplace_back("descending", tci.getDescending() ? "true" : "false"); } if (auto tfillpad = dyn_cast(op)) { - attrs.emplace_back( - "mode", pto::stringifyTFillPadMode(tfillpad.getMode()).str()); + auto kind = pto::inferTFillPadLoweringKindAfterMemoryPlanning(tfillpad); + if (failed(kind)) + return tfillpad.emitOpError( + "cannot infer a supported lowering; expand and in-place forms " + "require loc=vec, statically comparable physical shapes, and " + "resolved planned addresses"); + StringRef token; + switch (*kind) { + case pto::TFillPadLoweringKind::Normal: + token = "normal"; + break; + case pto::TFillPadLoweringKind::InPlace: + token = "in_place"; + break; + case pto::TFillPadLoweringKind::Expand: + token = "expand"; + break; + } + attrs.emplace_back("lowering_kind", token.str()); } if (auto tscatter = dyn_cast(op)) { if (auto maskPatternAttr = tscatter.getMaskPatternAttr()) { @@ -558,6 +576,7 @@ static void appendOpContextAttrs( op, attrs, pto::DivPrecision::HighPrecision) || tryAppendPrecisionType( op, attrs, pto::DivPrecision::HighPrecision)); + return success(); } static bool getStaticIntFromValue(Value value, int64_t &out) { @@ -774,21 +793,28 @@ static std::optional buildOperandTypeInfo(Value value) { return info; } -static std::optional buildSpecKey(Operation *op) { +static FailureOr buildSpecKey(Operation *op) { SpecKey key; key.opName = getTileOpName(op).str(); key.targetArch = getTargetArchString(op); for (unsigned i = 0; i < op->getNumOperands(); ++i) { auto info = buildOperandTypeInfo(op->getOperand(i)); - if (!info) - return std::nullopt; + if (!info) { + op->emitError("ExpandTileOp: cannot build specialization key for this " + "operand schema"); + return failure(); + } key.operands.push_back(*info); } - if (key.operands.empty()) - return std::nullopt; + if (key.operands.empty()) { + op->emitError( + "ExpandTileOp: cannot build a specialization key without operands"); + return failure(); + } - appendOpContextAttrs(op, key.contextAttrs); + if (failed(appendOpContextAttrs(op, key.contextAttrs))) + return failure(); return key; } @@ -1219,15 +1245,12 @@ LogicalResult ExpandState::expandTileOpsInFunction(func::FuncOp func, }); for (auto *op : tileOps) { - auto specKeyOpt = buildSpecKey(op); - if (!specKeyOpt) { - op->emitError( - "ExpandTileOp: cannot build specialization key for this operand schema"); + auto specKey = buildSpecKey(op); + if (failed(specKey)) return failure(); - } // Invoke the selected TileLib backend (with daemon-side caching). - func::FuncOp dslFn = invokeTileLib(*specKeyOpt, op, mod, ctx); + func::FuncOp dslFn = invokeTileLib(*specKey, op, mod, ctx); if (!dslFn) { StringRef opName = getTileOpName(op); op->emitError("ExpandTileOp: failed to instantiate TileLib template for " + diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index d5cad8a221..fb1713b630 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -11,6 +11,7 @@ #include "PTO/IR/PTOMultiBuffer.h" #include "PTO/IR/PTOTypeUtils.h" #include "PTO/Transforms/Passes.h" +#include "Utils.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Interfaces/ControlFlowInterfaces.h" @@ -270,8 +271,10 @@ static InplacePolicy getInplacePolicy(Operation *op) { "pto.tcolargmin", "pto.tcvt", "pto.txors", }); - if (auto fillPad = dyn_cast(op)) - policy.notInplaceSafe |= fillPad.getMode() == TFillPadMode::Expand; + if (auto fillPad = dyn_cast(op)) { + auto expanded = hasTFillPadExpandedPhysicalShape(fillPad); + policy.notInplaceSafe |= failed(expanded) || *expanded; + } if (name == "pto.tsel") { policy.forbidOutputAliasOperands.push_back(0); // mask diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index c473e1dd2d..a814e17c36 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -9778,16 +9778,16 @@ struct PTOInsertToEmitC : public OpConversionPattern { } }; -static StringRef getTFillPadModeToken(pto::TFillPadMode mode) { - switch (mode) { - case pto::TFillPadMode::Normal: +static StringRef getTFillPadModeToken(pto::TFillPadLoweringKind loweringKind) { + switch (loweringKind) { + case pto::TFillPadLoweringKind::Normal: return "pto::TFillPadMode::Normal"; - case pto::TFillPadMode::InPlace: + case pto::TFillPadLoweringKind::InPlace: return "pto::TFillPadMode::InPlace"; - case pto::TFillPadMode::Expand: + case pto::TFillPadLoweringKind::Expand: return "pto::TFillPadMode::Expand"; } - llvm_unreachable("unknown TFillPadMode"); + llvm_unreachable("unknown TFillPadLoweringKind"); } struct PTOFillPadToEmitC : public OpConversionPattern { @@ -9801,6 +9801,15 @@ struct PTOFillPadToEmitC : public OpConversionPattern { Value src = peelUnrealized(adaptor.getSrc()); Value dst = peelUnrealized(adaptor.getDst()); + auto loweringKind = pto::inferTFillPadLoweringKindAfterMemoryPlanning(op); + if (failed(loweringKind)) { + op.emitOpError( + "cannot infer a supported lowering; expand and in-place forms " + "require loc=vec, statically comparable physical shapes, and " + "resolved planned addresses"); + return failure(); + } + auto padValueTok = [&](pto::PadValue mode) -> StringRef { switch (mode) { case pto::PadValue::Null: @@ -9821,9 +9830,9 @@ struct PTOFillPadToEmitC : public OpConversionPattern { // tfillpad, so lowering can trust the preserved semantic contract. templateArgs = rewriter.getArrayAttr( {emitc::OpaqueAttr::get(ctx, padValueTok(padValueAttr.getValue()))}); - } else if (op.getMode() != pto::TFillPadMode::Normal) { + } else if (*loweringKind != pto::TFillPadLoweringKind::Normal) { templateArgs = rewriter.getArrayAttr( - {emitc::OpaqueAttr::get(ctx, getTFillPadModeToken(op.getMode()))}); + {emitc::OpaqueAttr::get(ctx, getTFillPadModeToken(*loweringKind))}); } rewriter.create( diff --git a/lib/PTO/Transforms/Utils.cpp b/lib/PTO/Transforms/Utils.cpp index 7165eb2c85..cc3f577d76 100644 --- a/lib/PTO/Transforms/Utils.cpp +++ b/lib/PTO/Transforms/Utils.cpp @@ -6,14 +6,17 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -#include "PTO/IR/PTO.h" #include "Utils.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/Support/ErrorHandling.h" +#include "PTO/IR/PTO.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/IR/Matchers.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/ErrorHandling.h" #define DEBUG_TYPE "pto-utils" #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") @@ -26,6 +29,96 @@ namespace pto { static constexpr llvm::StringLiteral kFrontendPipeIdAttrName = "__pto.frontend_id"; +FailureOr hasTFillPadExpandedPhysicalShape(TFillPadOp op) { + auto srcType = dyn_cast(op.getSrc().getType()); + auto dstType = dyn_cast(op.getDst().getType()); + if (!srcType || !dstType || srcType.getRank() != dstType.getRank()) + return failure(); + + bool expanded = false; + for (auto [srcDim, dstDim] : + llvm::zip_equal(srcType.getShape(), dstType.getShape())) { + if (srcDim == dstDim) + continue; + if (ShapedType::isDynamic(srcDim) || ShapedType::isDynamic(dstDim) || + dstDim < srcDim) + return failure(); + expanded = true; + } + return expanded; +} + +static Value peelTFillPadStorageAlias(Value value) { + constexpr unsigned kMaxDepth = 32; + for (unsigned depth = 0; value && depth < kMaxDepth; ++depth) { + Operation *def = value.getDefiningOp(); + if (!def) + break; + if (auto cast = dyn_cast(def)) { + if (cast.getNumOperands() != 1 || cast.getNumResults() != 1) + break; + value = cast.getOperand(0); + continue; + } + if (auto bitcast = dyn_cast(def)) { + value = bitcast.getSrc(); + continue; + } + if (auto reshape = dyn_cast(def)) { + value = reshape.getSrc(); + continue; + } + break; + } + return value; +} + +static bool haveSameKnownTFillPadStartAddress(Value src, Value dst) { + src = peelTFillPadStorageAlias(src); + dst = peelTFillPadStorageAlias(dst); + if (src == dst) + return true; + + auto srcAlloc = src.getDefiningOp(); + auto dstAlloc = dst.getDefiningOp(); + if (!srcAlloc || !dstAlloc || !srcAlloc.getAddr() || !dstAlloc.getAddr()) + return false; + + Value srcAddr = srcAlloc.getAddr(); + Value dstAddr = dstAlloc.getAddr(); + if (srcAddr == dstAddr) + return true; + + IntegerAttr srcConst; + IntegerAttr dstConst; + return matchPattern(srcAddr, m_Constant(&srcConst)) && + matchPattern(dstAddr, m_Constant(&dstConst)) && + srcConst.getValue() == dstConst.getValue(); +} + +FailureOr +inferTFillPadLoweringKindAfterMemoryPlanning(TFillPadOp op) { + FailureOr expanded = hasTFillPadExpandedPhysicalShape(op); + if (failed(expanded)) + return failure(); + + auto srcSpace = GetBufferSpaceAttr(op.getSrc()); + auto dstSpace = GetBufferSpaceAttr(op.getDst()); + bool isVec = srcSpace && dstSpace && + srcSpace->getAddressSpace() == AddressSpace::VEC && + dstSpace->getAddressSpace() == AddressSpace::VEC; + + if (*expanded) { + if (!isVec) + return failure(); + return TFillPadLoweringKind::Expand; + } + if (isVec && + haveSameKnownTFillPadStartAddress(op.getSrc(), op.getDst())) + return TFillPadLoweringKind::InPlace; + return TFillPadLoweringKind::Normal; +} + std::optional inferPhysicalSectionKindFromPipe(Operation *op) { auto pipeOp = dyn_cast_or_null(op); diff --git a/lib/PTO/Transforms/Utils.h b/lib/PTO/Transforms/Utils.h index 56f89b3e61..af341b4496 100644 --- a/lib/PTO/Transforms/Utils.h +++ b/lib/PTO/Transforms/Utils.h @@ -37,6 +37,12 @@ namespace mlir { namespace pto { + enum class TFillPadLoweringKind { + Normal, + InPlace, + Expand, + }; + enum class PhysicalSectionKind { Vector, Cube, @@ -45,6 +51,10 @@ namespace pto { std::optional inferPhysicalSectionKindFromPipe(Operation *op); + FailureOr hasTFillPadExpandedPhysicalShape(TFillPadOp op); + FailureOr + inferTFillPadLoweringKindAfterMemoryPlanning(TFillPadOp op); + const std::set LocalBufferSpace{ pto::AddressSpace::VEC, pto::AddressSpace::MAT, pto::AddressSpace::ACC, pto::AddressSpace::LEFT, pto::AddressSpace::RIGHT, pto::AddressSpace::BIAS, pto::AddressSpace::SCALING}; constexpr const uint8_t kBitsToByte = 8; diff --git a/lib/TileOps/a5/_fillpad.py b/lib/TileOps/a5/_fillpad.py index 0623c591cb..d7db069fac 100644 --- a/lib/TileOps/a5/_fillpad.py +++ b/lib/TileOps/a5/_fillpad.py @@ -203,16 +203,18 @@ def register_fillpad(): tags=("fillpad",), ) def template(src: pto.Tile, dst: pto.Tile): - mode = pto.get_op_attr("mode", "normal") + lowering_kind = pto.get_op_attr("lowering_kind", "normal") src_valid_rows, src_valid_cols = src.valid_shape dst_valid_rows, dst_valid_cols = dst.valid_shape lanes = pto.elements_per_vreg(dst.dtype) aligned_cols = (src_valid_cols // lanes) * lanes - if mode == "in_place": + if lowering_kind == "in_place": _fill_inplace(dst, src_valid_rows, src_valid_cols, dst_valid_rows, dst_valid_cols) return _copy_region(src, dst, src_valid_rows, 0, aligned_cols) - fill_row_stop = dst_valid_rows if mode == "expand" else src_valid_rows + fill_row_stop = ( + dst_valid_rows if lowering_kind == "expand" else src_valid_rows + ) scalar_tail_start = _scalar_tail_start(dst, lanes) _fill(dst, 0, fill_row_stop, aligned_cols, dst_valid_cols, scalar_tail_start=scalar_tail_start) _copy_region(src, dst, src_valid_rows, aligned_cols, src_valid_cols) diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 5c45b1846d..0327b9be7a 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -4150,31 +4150,11 @@ def tpartmin(src0, src1, dst): ) -def _tfillpad_mode_attr(mode): - if isinstance(mode, Attribute): - return mode - if isinstance(mode, str): - token = mode.strip().lower().replace("-", "_") - aliases = { - "normal": _pto.TFillPadMode.Normal, - "inplace": _pto.TFillPadMode.InPlace, - "in_place": _pto.TFillPadMode.InPlace, - "expand": _pto.TFillPadMode.Expand, - } - if token not in aliases: - raise ValueError( - "tfillpad mode must be 'normal', 'in_place', or 'expand'" - ) - mode = aliases[token] - return _pto.TFillPadModeAttr.get(mode) - - -def tfillpad(src, dst, *, mode="normal"): - """``pto.tfillpad ins(src) outs(dst)`` with an explicit ISA mode.""" +def tfillpad(src, dst): + """``pto.tfillpad ins(src) outs(dst)`` with compiler-inferred lowering.""" _pto.tfillpad( unwrap_surface_value(src), unwrap_surface_value(dst), - mode=_tfillpad_mode_attr(mode), ) diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index 1ce5bf4903..d4f1774b5e 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -134,22 +134,19 @@ def test_tile_partial_and_fillpad_names_are_exposed_without_legacy_names(self): with self.subTest(name=name): self.assertFalse(hasattr(pto.tile, name), name) - def test_tile_fillpad_dispatches_one_op_with_mode(self): + def test_tile_fillpad_dispatches_one_op_without_mode(self): src = object() dst = object() - mode_attr = object() with patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ - patch.object(_ops, "_tfillpad_mode_attr", return_value=mode_attr) as build_mode, \ patch.object(_ops._pto, "tfillpad") as tfillpad: - pto.tile.fillpad(src, dst, mode="expand") + pto.tile.fillpad(src, dst) - build_mode.assert_called_once_with("expand") - tfillpad.assert_called_once_with(src, dst, mode=mode_attr) + tfillpad.assert_called_once_with(src, dst) - def test_tile_fillpad_rejects_unknown_mode(self): - with self.assertRaisesRegex(ValueError, "normal.*in_place.*expand"): - _ops._tfillpad_mode_attr("automatic") + def test_tile_fillpad_rejects_mode_argument(self): + with self.assertRaises(TypeError): + pto.tile.fillpad(object(), object(), mode="expand") def test_sync_flag_names_are_exposed_without_legacy_aliases(self): preferred_names = [ diff --git a/python/pto/dialects/pto.py b/python/pto/dialects/pto.py index 2c868b9bd4..161cce5404 100644 --- a/python/pto/dialects/pto.py +++ b/python/pto/dialects/pto.py @@ -73,8 +73,6 @@ def _export_optional_cext_symbol(name): PadValueAttr = _pto_mod.PadValueAttr CompactMode = _pto_mod.CompactMode CompactModeAttr = _pto_mod.CompactModeAttr -TFillPadMode = _pto_mod.TFillPadMode -TFillPadModeAttr = _pto_mod.TFillPadModeAttr AccToVecMode = _pto_mod.AccToVecMode AccToVecModeAttr = _pto_mod.AccToVecModeAttr TInsertMode = _pto_mod.TInsertMode @@ -254,8 +252,6 @@ def fence_scope_attr_builder(value, context=None): "PadValueAttr", "CompactMode", "CompactModeAttr", - "TFillPadMode", - "TFillPadModeAttr", "AccToVecMode", "AccToVecModeAttr", "TInsertMode", diff --git a/test/lit/pto/fillpad_tile_native.pto b/test/lit/pto/fillpad_tile_native.pto index 434bb11a4a..70f701ed14 100644 --- a/test/lit/pto/fillpad_tile_native.pto +++ b/test/lit/pto/fillpad_tile_native.pto @@ -14,9 +14,10 @@ module { func.func private @tfillpad_arg( - %tile: !pto.tile_buf) { - pto.tfillpad ins(%tile : !pto.tile_buf) - outs(%tile : !pto.tile_buf) + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tfillpad ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) return } @@ -24,16 +25,17 @@ module { %tile: !pto.tile_buf) { pto.tfillpad ins(%tile : !pto.tile_buf) outs(%tile : !pto.tile_buf) - {mode = #pto.tfillpad_mode} return } } // NATIVE-LABEL: func.func private @tfillpad_arg( // NATIVE: pto.tfillpad ins(%arg0 +// NATIVE-SAME: outs(%arg1 // NATIVE-LABEL: func.func private @tfillpad_inplace_arg( // NATIVE: pto.tfillpad ins(%arg0 -// NATIVE-SAME: mode = #pto.tfillpad_mode +// NATIVE-SAME: outs(%arg0 +// NATIVE-NOT: mode // NATIVE-NOT: memref< // EMITC-LABEL: tfillpad_arg( diff --git a/test/lit/pto/movement_metadata_tile_native.pto b/test/lit/pto/movement_metadata_tile_native.pto index f8ee39c914..66d96f09ee 100644 --- a/test/lit/pto/movement_metadata_tile_native.pto +++ b/test/lit/pto/movement_metadata_tile_native.pto @@ -11,7 +11,7 @@ module { func.func private @tfillpad_expand_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { - pto.tfillpad ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) {mode = #pto.tfillpad_mode} + pto.tfillpad ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) return } func.func private @tget_scale_addr_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { @@ -22,7 +22,7 @@ module { // NATIVE-LABEL: @tfillpad_expand_arg // NATIVE: pto.tfillpad -// NATIVE-SAME: mode = #pto.tfillpad_mode +// NATIVE-NOT: mode // NATIVE-LABEL: @tget_scale_addr_arg // NATIVE: pto.tget_scale_addr // NATIVE-NOT: memref< diff --git a/test/lit/pto/tfillpad_inplace_alias_lowering.pto b/test/lit/pto/tfillpad_inplace_alias_lowering.pto index e5a059766c..d8ddc62cc4 100644 --- a/test/lit/pto/tfillpad_inplace_alias_lowering.pto +++ b/test/lit/pto/tfillpad_inplace_alias_lowering.pto @@ -1,14 +1,28 @@ -// RUN: ptoas %s | FileCheck %s +// RUN: ptoas --pto-level=level3 %s | FileCheck %s module { func.func @tfillpad_inplace_alias() { - %tile = pto.alloc_tile : !pto.tile_buf - pto.tfillpad ins(%tile : !pto.tile_buf) - outs(%tile : !pto.tile_buf) - {mode = #pto.tfillpad_mode} + %c0 = arith.constant 0 : i64 + %src = pto.alloc_tile addr = %c0 : !pto.tile_buf + %dst = pto.alloc_tile addr = %c0 : !pto.tile_buf + pto.tfillpad ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @tfillpad_different_addresses() { + %c0 = arith.constant 0 : i64 + %c4096 = arith.constant 4096 : i64 + %src = pto.alloc_tile addr = %c0 : !pto.tile_buf + %dst = pto.alloc_tile addr = %c4096 : !pto.tile_buf + pto.tfillpad ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) return } } // CHECK-LABEL: AICORE void tfillpad_inplace_alias( // CHECK: TFILLPAD( +// CHECK-LABEL: AICORE void tfillpad_different_addresses( +// CHECK-NOT: TFillPadMode::InPlace +// CHECK: TFILLPAD( diff --git a/test/lit/pto/tfillpad_non_normal_mat_invalid.pto b/test/lit/pto/tfillpad_non_normal_mat_invalid.pto index 8b4b52e454..9c42d41ec8 100644 --- a/test/lit/pto/tfillpad_non_normal_mat_invalid.pto +++ b/test/lit/pto/tfillpad_non_normal_mat_invalid.pto @@ -1,14 +1,14 @@ // RUN: not ptoas --pto-arch=a3 %s -o /dev/null 2>&1 | FileCheck %s module { - func.func @tfillpad_expand_mat_invalid( + func.func @tfillpad_explicit_mode_invalid( %src: !pto.tile_buf, %dst: !pto.tile_buf) { pto.tfillpad ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) - {mode = #pto.tfillpad_mode} + {mode = "expand"} return } } -// CHECK: error: expects non-normal TFILLPAD mode only for loc=vec +// CHECK: error: 'pto.tfillpad' op does not accept 'mode'; PTOAS infers TFILLPAD lowering from physical shape and planned addresses diff --git a/test/lit/pto/tfillpad_plan_memory_inference.pto b/test/lit/pto/tfillpad_plan_memory_inference.pto new file mode 100644 index 0000000000..346d6e083f --- /dev/null +++ b/test/lit/pto/tfillpad_plan_memory_inference.pto @@ -0,0 +1,22 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --plan-memory-impl=modern %s | FileCheck %s + +module { + func.func @tfillpad_plan_memory_alias() { + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tfillpad ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK-LABEL: AICORE void tfillpad_plan_memory_alias( +// CHECK: TFILLPAD( diff --git a/test/lit/pto/tfillpad_same_ssa_lowers_to_tfillpad.pto b/test/lit/pto/tfillpad_same_ssa_lowers_to_tfillpad.pto index 80dd84f991..be9a968200 100644 --- a/test/lit/pto/tfillpad_same_ssa_lowers_to_tfillpad.pto +++ b/test/lit/pto/tfillpad_same_ssa_lowers_to_tfillpad.pto @@ -10,5 +10,4 @@ module { } // CHECK-LABEL: AICORE void tfillpad_same_ssa( -// CHECK: TFILLPAD( -// CHECK-NOT: TFillPadMode::InPlace +// CHECK: TFILLPAD( diff --git a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_expand.pto b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_expand.pto index 4e0d315640..97b3135360 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_expand.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_expand.pto @@ -7,14 +7,14 @@ // See LICENSE in the root of the software repository for the full text of the License. // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline -// expands pto.tfillpad in expand mode via the PTODSL TileLib template +// expands pto.tfillpad with compiler-inferred expand lowering via the PTODSL TileLib template // // Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s // After the full tile-op-expand path on the VPTO backend, the original -// pto.tfillpad in expand mode should be lowered to vector-style VPTO IR. +// Compiler-inferred expand lowering should produce vector-style VPTO IR. // CHECK: func.func @TFILLPAD_EXPAND // CHECK-NOT: pto.tfillpad ins // CHECK: pto.vecscope @@ -38,7 +38,6 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { blayout=row_major, slayout=none_box, fractal=512, pad=0>) outs(%dst : !pto.tile_buf) - {mode = #pto.tfillpad_mode} return } } diff --git a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto index 3b1adfb019..7a27d14dd2 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto @@ -7,14 +7,14 @@ // See LICENSE in the root of the software repository for the full text of the License. // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline -// expands pto.tfillpad (inplace mode) via the PTODSL TileLib template +// expands pto.tfillpad with compiler-inferred in-place lowering via the PTODSL TileLib template // // Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s +// RUN: ptoas --pto-level=level3 --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s // After the full tile-op-expand path on the VPTO backend, the original -// pto.tfillpad (inplace) should be lowered to vector-style VPTO IR. +// Compiler-inferred in-place lowering should produce vector-style VPTO IR. // CHECK: func.func @TFILLPAD_INPLACE // CHECK-NOT: pto.tfillpad ins // CHECK: %[[MAX:.*]] = arith.constant 3.40282347E+38 : f32 @@ -28,10 +28,11 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { func.func @TFILLPAD_INPLACE() { - %src = pto.alloc_tile + %c0 = arith.constant 0 : i64 + %src = pto.alloc_tile addr = %c0 : !pto.tile_buf - %dst = pto.alloc_tile + %dst = pto.alloc_tile addr = %c0 : !pto.tile_buf @@ -41,7 +42,6 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { blayout=row_major, slayout=none_box, fractal=512, pad=3>) outs(%dst : !pto.tile_buf) - {mode = #pto.tfillpad_mode} return } } diff --git a/test/samples/Fillpad/fillpad_expand.py b/test/samples/Fillpad/fillpad_expand.py index e064122ae8..89751ec145 100644 --- a/test/samples/Fillpad/fillpad_expand.py +++ b/test/samples/Fillpad/fillpad_expand.py @@ -28,8 +28,6 @@ def build(): bl = pto.BLayoutAttr.get(pto.BLayout.RowMajor, ctx) sl = pto.SLayoutAttr.get(pto.SLayout.NoneBox, ctx) pd = pto.PadValueAttr.get(pto.PadValue.Zero, ctx) - mode = pto.TFillPadModeAttr.get(pto.TFillPadMode.Expand, ctx) - fractal_ab_size = pto.TileConfig.fractalABSize cfg = pto.TileBufConfigAttr.get(bl, sl, fractal_ab_size, pd, ctx) tile_buf_32_16 = pto.TileBufType.get([32, 16], f32, vec, [32, 16], cfg, ctx) @@ -63,7 +61,7 @@ def build(): dst_tb = pto.AllocTileOp(tile_buf_32_32).result pto.TLoadOp(None, src_sv, src_tb) - pto.TFillPadOp(src_tb, dst_tb, mode=mode) + pto.TFillPadOp(src_tb, dst_tb) pto.TStoreOp(None, dst_tb, dst_sv) func.ReturnOp([]) diff --git a/test/samples/Fillpad/fillpad_expand_invalid.py b/test/samples/Fillpad/fillpad_expand_invalid.py index 2f17f07375..8b32e800b6 100644 --- a/test/samples/Fillpad/fillpad_expand_invalid.py +++ b/test/samples/Fillpad/fillpad_expand_invalid.py @@ -23,8 +23,6 @@ def build(): bl = pto.BLayoutAttr.get(pto.BLayout.RowMajor, ctx) sl = pto.SLayoutAttr.get(pto.SLayout.NoneBox, ctx) pd = pto.PadValueAttr.get(pto.PadValue.Zero, ctx) - mode = pto.TFillPadModeAttr.get(pto.TFillPadMode.Expand, ctx) - fractal_ab_size = pto.TileConfig.fractalABSize cfg = pto.TileBufConfigAttr.get(bl, sl, fractal_ab_size, pd, ctx) src_ty = pto.TileBufType.get([32, 32], f32, vec, [32, 32], cfg, ctx) @@ -39,7 +37,7 @@ def build(): with InsertionPoint(entry): src = pto.AllocTileOp(src_ty).result dst = pto.AllocTileOp(dst_ty).result - pto.TFillPadOp(src, dst, mode=mode) + pto.TFillPadOp(src, dst) func.ReturnOp([]) ok = m.operation.verify() diff --git a/test/samples/Fillpad/fillpad_expand_pad_null_invalid.py b/test/samples/Fillpad/fillpad_expand_pad_null_invalid.py index d194a093dc..ea6fa854af 100644 --- a/test/samples/Fillpad/fillpad_expand_pad_null_invalid.py +++ b/test/samples/Fillpad/fillpad_expand_pad_null_invalid.py @@ -24,8 +24,6 @@ def build(): sl = pto.SLayoutAttr.get(pto.SLayout.NoneBox, ctx) src_pd = pto.PadValueAttr.get(pto.PadValue.Zero, ctx) dst_pd = pto.PadValueAttr.get(pto.PadValue.Null, ctx) - mode = pto.TFillPadModeAttr.get(pto.TFillPadMode.Expand, ctx) - fractal_ab_size = pto.TileConfig.fractalABSize src_cfg = pto.TileBufConfigAttr.get(bl, sl, fractal_ab_size, src_pd, ctx) dst_cfg = pto.TileBufConfigAttr.get(bl, sl, fractal_ab_size, dst_pd, ctx) @@ -41,7 +39,7 @@ def build(): with InsertionPoint(entry): src = pto.AllocTileOp(src_ty).result dst = pto.AllocTileOp(dst_ty).result - pto.TFillPadOp(src, dst, mode=mode) + pto.TFillPadOp(src, dst) func.ReturnOp([]) ok = m.operation.verify() diff --git a/test/samples/Fillpad/fillpad_inplace.py b/test/samples/Fillpad/fillpad_inplace.py index e5f0dffd5c..48dd629ff9 100644 --- a/test/samples/Fillpad/fillpad_inplace.py +++ b/test/samples/Fillpad/fillpad_inplace.py @@ -27,8 +27,6 @@ def build(): bl = pto.BLayoutAttr.get(pto.BLayout.RowMajor, ctx) sl = pto.SLayoutAttr.get(pto.SLayout.NoneBox, ctx) pd = pto.PadValueAttr.get(pto.PadValue.Zero, ctx) - mode = pto.TFillPadModeAttr.get(pto.TFillPadMode.InPlace, ctx) - cfg = pto.TileBufConfigAttr.get(bl, sl, pto.TileConfig.fractalABSize, pd, ctx) tile_ty = pto.TileBufType.get([32, 32], f32, vec, [32, 32], cfg, ctx) @@ -51,7 +49,7 @@ def build(): tile = pto.AllocTileOp(tile_ty).result pto.TLoadOp(None, sv0, tile) - pto.TFillPadOp(tile, tile, mode=mode) + pto.TFillPadOp(tile, tile) pto.TStoreOp(None, tile, sv1) func.ReturnOp([]) diff --git a/test/samples/runop.sh b/test/samples/runop.sh index 2db16abeda..d8258a79c6 100755 --- a/test/samples/runop.sh +++ b/test/samples/runop.sh @@ -988,13 +988,8 @@ PY fi if [[ "$base" == "fillpad" ]]; then - if ! grep -Fq "TFILLPAD(" "$cpp"; then - echo -e "${A}(${base}.py)\tFAIL\tmissing TFILLPAD() lowering for pto.tfillpad" - overall=1 - continue - fi - if grep -Fq "TFillPadMode::" "$cpp"; then - echo -e "${A}(${base}.py)\tFAIL\tnormal pto.tfillpad should use the default ISA mode" + if ! grep -Fq "TFILLPAD" "$cpp"; then + echo -e "${A}(${base}.py)\tFAIL\tmissing compiler-inferred TFILLPAD lowering" overall=1 continue fi diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_expand/tfillpad_expand.pto b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_expand/tfillpad_expand.pto index 3de41b5a0a..f00126f112 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_expand/tfillpad_expand.pto +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_expand/tfillpad_expand.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// TileLang ST kernels for pto.tfillpad in expand mode: copy src to dst and fill padding. +// TileLang ST kernels for compiler-inferred tfillpad expand lowering: copy src to dst and fill padding. // Matches C++ test cases: case 8, 9 // Compiled by ptoas --enable-insert-sync --enable-tile-op-expand --pto-backend=vpto // @@ -49,16 +49,15 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind -> !pto.partition_tensor_view<1x1x1x64x16xui16> %src = pto.alloc_tile - : !pto.tile_buf + : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x63x7xui16>) - outs(%src : !pto.tile_buf) + outs(%src : !pto.tile_buf) - pto.tfillpad ins(%src : !pto.tile_buf) + pto.tfillpad ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) - {mode = #pto.tfillpad_mode} pto.tstore ins(%dst : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x64x16xui16>) @@ -101,17 +100,16 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind + : !pto.tile_buf // Dst tile: FillPadVal=Max (pad=2), dst physical=260x32, v_row=260, v_col=32 (full output) %dst = pto.alloc_tile : !pto.tile_buf pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x259x7xui16>) - outs(%src : !pto.tile_buf) + outs(%src : !pto.tile_buf) - pto.tfillpad ins(%src : !pto.tile_buf) + pto.tfillpad ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) - {mode = #pto.tfillpad_mode} pto.tstore ins(%dst : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x32xui16>) diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto index 1c2b40183b..aae68a26e5 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tfillpad_inplace/tfillpad_inplace.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// TileLang ST kernels for pto.tfillpad (inplace mode). +// TileLang ST kernels for compiler-inferred pto.tfillpad in-place lowering. // Matches C++ reference test case: Case 5 // Compiled by ptoas --enable-insert-sync --enable-tile-op-expand --pto-backend=vpto // @@ -58,7 +58,6 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) outs(%dst_tile : !pto.tile_buf) - {mode = #pto.tfillpad_mode} pto.tstore ins(%dst_tile : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x64x16xf32>) @@ -110,7 +109,6 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) outs(%dst_tile : !pto.tile_buf) - {mode = #pto.tfillpad_mode} pto.tstore ins(%dst_tile : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x16xf32>) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_expand/tfillpad_expand.pto b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_expand/tfillpad_expand.pto index 3d68127035..329cb95a70 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_expand/tfillpad_expand.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_expand/tfillpad_expand.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// TileLang ST kernels for pto.tfillpad in expand mode: copy src to dst and fill padding. +// TileLang ST kernels for compiler-inferred tfillpad expand lowering: copy src to dst and fill padding. // Matches C++ test cases: case 8, 9 // Compiled by ptoas --enable-insert-sync --enable-tile-op-expand --pto-backend=vpto // @@ -50,19 +50,18 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind -> !pto.partition_tensor_view<1x1x1x260x32xui16> - // Src tile: LoadPadVal=Min (pad=3), src physical=260x32, v_row=259, v_col=7 + // Src tile: LoadPadVal=Min (pad=3), src physical=259x32, v_row=259, v_col=7 %src = pto.alloc_tile - : !pto.tile_buf + : !pto.tile_buf // Dst tile: FillPadVal=Max (pad=2), dst physical=260x32, v_row=260, v_col=32 (full output) %dst = pto.alloc_tile : !pto.tile_buf pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x259x7xui16>) - outs(%src : !pto.tile_buf) + outs(%src : !pto.tile_buf) - pto.tfillpad ins(%src : !pto.tile_buf) + pto.tfillpad ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) - {mode = #pto.tfillpad_mode} pto.tstore ins(%dst : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x32xui16>) @@ -102,19 +101,18 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind -> !pto.partition_tensor_view<1x1x1x260x64xi8> - // Src tile: LoadPadVal=Min (pad=3), src physical=260x64, v_row=259, v_col=7 + // Src tile: LoadPadVal=Min (pad=3), src physical=259x64, v_row=259, v_col=7 %src = pto.alloc_tile - : !pto.tile_buf + : !pto.tile_buf // Dst tile: FillPadVal=Max (pad=2), dst physical=260x64, v_row=260, v_col=64 (full output) %dst = pto.alloc_tile : !pto.tile_buf pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x259x7xi8>) - outs(%src : !pto.tile_buf) + outs(%src : !pto.tile_buf) - pto.tfillpad ins(%src : !pto.tile_buf) + pto.tfillpad ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) - {mode = #pto.tfillpad_mode} pto.tstore ins(%dst : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x64xi8>) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto index 32e797a038..5965c0795c 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tfillpad_inplace/tfillpad_inplace.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// TileLang ST kernels for pto.tfillpad (inplace mode). +// TileLang ST kernels for compiler-inferred pto.tfillpad in-place lowering. // Matches C++ reference test case: Case 5 // Compiled by ptoas --enable-insert-sync --enable-tile-op-expand --pto-backend=vpto // @@ -63,7 +63,6 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) outs(%dst_tile : !pto.tile_buf) - {mode = #pto.tfillpad_mode} pto.tstore ins(%dst_tile : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x260x16xf32>) From db3c90d3e902647e64d1d056ac510f041f266628 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Thu, 6 Aug 2026 13:08:30 +0800 Subject: [PATCH 019/122] fix: expand A5 tinsert fp tile templates --- lib/TileOps/a5/tinsert.py | 250 ++++++++++++++++++ ...nd_tile_op_tilelang_tinsert_fp_acc2mat.pto | 58 ++++ 2 files changed, 308 insertions(+) create mode 100644 test/lit/vpto/expand_tile_op_tilelang_tinsert_fp_acc2mat.pto diff --git a/lib/TileOps/a5/tinsert.py b/lib/TileOps/a5/tinsert.py index b01f075f7d..c0fab9d7cd 100644 --- a/lib/TileOps/a5/tinsert.py +++ b/lib/TileOps/a5/tinsert.py @@ -80,6 +80,119 @@ def _vec_to_vec_nd_scalar(src_memory_space, dst_memory_space, src_config, dst_co _DTYPES = [(dtype, "i32", "i32", dtype) for dtype in NUMERIC_DTYPES] +_FP_DTYPES = ( + ("f32", "i32", "i32", "i8", "f32"), + ("f32", "i32", "i32", "si8", "f32"), + ("f32", "i32", "i32", "ui8", "f32"), + ("f32", "i32", "i32", "f16", "f32"), + ("f32", "i32", "i32", "bf16", "f32"), + ("f32", "i32", "i32", "f32", "f32"), + ("i32", "i32", "i32", "i8", "f32"), + ("si32", "i32", "i32", "i8", "f32"), + ("i32", "i32", "i32", "si8", "f32"), + ("si32", "i32", "i32", "si8", "f32"), + ("i32", "i32", "i32", "ui8", "f32"), + ("si32", "i32", "i32", "ui8", "f32"), + ("i32", "i32", "i32", "f16", "f32"), + ("si32", "i32", "i32", "f16", "f32"), + ("i32", "i32", "i32", "bf16", "f32"), + ("si32", "i32", "i32", "bf16", "f32"), +) + + +def _tinsert_fp_quant_mode(src_dtype, dst_dtype): + modes = { + ("f32", "i8"): "qf322b8_pre_vec", + ("f32", "si8"): "qf322b8_pre_vec", + ("f32", "ui8"): "qf322b8_pre_vec", + ("f32", "f16"): "qf322f16_pre_vec", + ("f32", "bf16"): "qf322bf16_pre_vec", + ("f32", "f32"): "qf322f32_pre_vec", + ("i32", "i8"): "req8_vec", + ("si32", "i8"): "req8_vec", + ("i32", "si8"): "req8_vec", + ("si32", "si8"): "req8_vec", + ("i32", "ui8"): "req8_vec", + ("si32", "ui8"): "req8_vec", + ("i32", "f16"): "deqf16_vec", + ("si32", "f16"): "deqf16_vec", + ("i32", "bf16"): "qs322bf16_pre_vec", + ("si32", "bf16"): "qs322bf16_pre_vec", + } + return modes[(str(src_dtype), str(dst_dtype))] + + +def _acc_to_mat_fp( + src_kind, + dst_kind, + src_memory_space, + dst_memory_space, + src_config, + dst_config, + fp_kind, + fp_memory_space, + **_, +): + return ( + _acc_to_mat( + src_kind, + dst_kind, + src_memory_space, + dst_memory_space, + src_config, + dst_config, + ) + and fp_kind == "tile" + and fp_memory_space == "scaling" + ) + + +def _acc_to_vec_nd_fp( + src_kind, + dst_kind, + src_memory_space, + dst_memory_space, + dst_config, + fp_kind, + fp_memory_space, + **_, +): + return ( + _acc_to_vec_nd( + src_kind, + dst_kind, + src_memory_space, + dst_memory_space, + dst_config, + ) + and fp_kind == "tile" + and fp_memory_space == "scaling" + ) + + +def _acc_to_vec_nz_fp( + src_kind, + dst_kind, + src_memory_space, + dst_memory_space, + dst_config, + fp_kind, + fp_memory_space, + **_, +): + return ( + _acc_to_vec_nz( + src_kind, + dst_kind, + src_memory_space, + dst_memory_space, + dst_config, + ) + and fp_kind == "tile" + and fp_memory_space == "scaling" + ) + + def _acc_to_vec_store_kwargs(): dst_mode = 0 split_mode = None @@ -142,6 +255,47 @@ def template_tinsert_acc_to_mat_basic( ) +@tilelib.tile_template( + op="pto.tinsert", + target="a5", + name="template_tinsert_fp_acc_to_mat", + dtypes=_FP_DTYPES, + iteration_axis="none", + op_engine="other", + op_class="movement", + constraints=[_acc_to_mat_fp], + priority=1, + id=5, + loop_depth=0, + is_post_update=False, + tags=("insert", "acc", "mat", "fp"), +) +def template_tinsert_fp_acc_to_mat( + src: pto.Tile, + index_row: pto.i32, + index_col: pto.i32, + dst: pto.Tile, + fp: pto.Tile, +): + elem_bytes = pto.bytewidth(dst.dtype) + c0_size = BLOCK_BYTE_SIZE // elem_bytes + valid_rows, valid_cols = src.valid_shape + n_size = (valid_cols + c0_size - 1) // c0_size * c0_size + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst.shape[0] * c0_size * col_block + index_row * c0_size + col_mod + + pto.mte_l0c_l1( + src.as_ptr(), + pto.addptr(dst.as_ptr(), dst_offset), + valid_rows, + n_size, + src.shape[0] * pto.bytewidth(src.dtype), + dst.shape[0] * c0_size * elem_bytes, + pre_quant=(fp.as_ptr(), _tinsert_fp_quant_mode(src.dtype, dst.dtype)), + ) + + @tilelib.tile_template( op="pto.tinsert", target="a5", @@ -213,6 +367,49 @@ def template_tinsert_acc_to_vec_nd_basic( ) +@tilelib.tile_template( + op="pto.tinsert", + target="a5", + name="template_tinsert_fp_acc_to_vec_nd", + dtypes=_FP_DTYPES, + iteration_axis="none", + op_engine="other", + op_class="movement", + constraints=[_acc_to_vec_nd_fp], + priority=1, + id=6, + loop_depth=0, + is_post_update=False, + tags=("insert", "acc", "vec", "nd", "fp"), +) +def template_tinsert_fp_acc_to_vec_nd( + src: pto.Tile, + index_row: pto.i32, + index_col: pto.i32, + dst: pto.Tile, + fp: pto.Tile, +): + elem_bytes = pto.bytewidth(dst.dtype) + c0_size = BLOCK_BYTE_SIZE // elem_bytes + valid_rows, valid_cols_raw = src.valid_shape + valid_cols = (valid_cols_raw + c0_size - 1) // c0_size * c0_size + dst_ptr = pto.addptr(dst.as_ptr(), index_row * dst.shape[1] + index_col) + dst_mode, kwargs = _acc_to_vec_store_kwargs() + kwargs["layout"] = "nz2nd" + + pto.mte_l0c_ub( + src.as_ptr(), + dst_ptr, + valid_rows, + valid_cols, + (valid_rows + 15) // 16 * 16, + dst.shape[1], + dst_mode, + pre_quant=(fp.as_ptr(), _tinsert_fp_quant_mode(src.dtype, dst.dtype)), + **kwargs, + ) + + @tilelib.tile_template( op="pto.tinsert", target="a5", @@ -263,6 +460,59 @@ def template_tinsert_acc_to_vec_nz_basic( ) +@tilelib.tile_template( + op="pto.tinsert", + target="a5", + name="template_tinsert_fp_acc_to_vec_nz", + dtypes=_FP_DTYPES, + iteration_axis="none", + op_engine="other", + op_class="movement", + constraints=[_acc_to_vec_nz_fp], + priority=1, + id=7, + loop_depth=0, + is_post_update=False, + tags=("insert", "acc", "vec", "nz", "fp"), +) +def template_tinsert_fp_acc_to_vec_nz( + src: pto.Tile, + index_row: pto.i32, + index_col: pto.i32, + dst: pto.Tile, + fp: pto.Tile, +): + elem_bytes = pto.bytewidth(dst.dtype) + c0_size = BLOCK_BYTE_SIZE // elem_bytes + valid_rows, valid_cols_raw = src.valid_shape + valid_cols_align = 16 if str(dst.dtype) == "f32" else c0_size + valid_cols = ( + (valid_cols_raw + valid_cols_align - 1) + // valid_cols_align + * valid_cols_align + ) + + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = ( + dst.shape[0] * c0_size * col_block + index_row * c0_size + col_mod + ) + dst_mode, kwargs = _acc_to_vec_store_kwargs() + kwargs["layout"] = ("nz2nz", 0) + + pto.mte_l0c_ub( + src.as_ptr(), + pto.addptr(dst.as_ptr(), dst_offset), + valid_rows, + valid_cols, + (valid_rows + 15) // 16 * 16 * pto.bytewidth(src.dtype), + dst.shape[0] * c0_size * elem_bytes, + dst_mode, + pre_quant=(fp.as_ptr(), _tinsert_fp_quant_mode(src.dtype, dst.dtype)), + **kwargs, + ) + + @tilelib.tile_template( op="pto.tinsert", target="a5", diff --git a/test/lit/vpto/expand_tile_op_tilelang_tinsert_fp_acc2mat.pto b/test/lit/vpto/expand_tile_op_tilelang_tinsert_fp_acc2mat.pto new file mode 100644 index 0000000000..02b828890c --- /dev/null +++ b/test/lit/vpto/expand_tile_op_tilelang_tinsert_fp_acc2mat.pto @@ -0,0 +1,58 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// The unified five-operand pto.tinsert fp form must be discoverable for every +// A5 Acc destination layout, not only by direct EmitC lowering. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s + +// CHECK-LABEL: func.func @TINSERT_FP_ACC_TO_MAT() +// CHECK-NOT: pto.tinsert +// CHECK: pto.set_fpc +// CHECK: pto.copy_matrix_cc_to_cbuf + +// CHECK-LABEL: func.func @TINSERT_FP_ACC_TO_VEC_ND() +// CHECK-NOT: pto.tinsert +// CHECK: pto.set_fpc +// CHECK: pto.copy_matrix_cc_to_ub + +// CHECK-LABEL: func.func @TINSERT_FP_ACC_TO_VEC_NZ() +// CHECK-NOT: pto.tinsert +// CHECK: pto.set_fpc +// CHECK: pto.copy_matrix_cc_to_ub + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @TINSERT_FP_ACC_TO_MAT() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @TINSERT_FP_ACC_TO_VEC_ND() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @TINSERT_FP_ACC_TO_VEC_NZ() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} From 56eee628c04078c19818d4900a62703b335f7230 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 13:47:34 +0800 Subject: [PATCH 020/122] fix(ptodsl): merge section-local conditional bindings (#1162) --- ptodsl/ptodsl/_ast_rewrite.py | 5 +++++ ptodsl/tests/test_section.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index 8a12ae7f57..c07e863e42 100644 --- a/ptodsl/ptodsl/_ast_rewrite.py +++ b/ptodsl/ptodsl/_ast_rewrite.py @@ -200,6 +200,11 @@ def visit_For(self, node): def visit_If(self, node): node.test = self.visit(node.test) + # Both branches of a runtime conditional share one authored binding. + # Reserve its section-local alias before visiting either branch so the + # branch merge does not treat the second branch as a new binding. + common_targets = _name_info(node.body).stores & _name_info(node.orelse).stores + self._activate_targets(common_targets) entry_env = dict(self._env) node.body, body_env = self._visit_block(node.body, entry_env) node.orelse, else_env = self._visit_block(node.orelse, entry_env) diff --git a/ptodsl/tests/test_section.py b/ptodsl/tests/test_section.py index 15ed6732a6..60c993cccb 100644 --- a/ptodsl/tests/test_section.py +++ b/ptodsl/tests/test_section.py @@ -132,6 +132,32 @@ def lexical_section_conditional_rebinding_probe(): pto.wait_flag("S", "MTE2", event_id=value) +@pto.jit(target="a5", mode="explicit") +def lexical_section_sibling_conditional_rebinding_probe(): + one = pto.const(1, dtype=pto.i32) + two = pto.const(2, dtype=pto.i32) + m_tile = pto.const(0, dtype=pto.i32) + n_tile = pto.const(0, dtype=pto.i32) + with pto.section("cube"): + if pto.get_block_idx() < one: + m_tile = one + n_tile = one + else: + m_tile = two + n_tile = two + pto.wait_flag("S", "MTE2", event_id=m_tile) + pto.wait_flag("S", "MTE2", event_id=n_tile) + with pto.section("vector"): + if pto.get_block_idx() < one: + m_tile = one + n_tile = one + else: + m_tile = two + n_tile = two + pto.wait_flag("MTE2", "S", event_id=m_tile) + pto.wait_flag("MTE2", "S", event_id=n_tile) + + @pto.jit(target="a5", mode="explicit") def lexical_section_loop_carry_probe(): one = pto.const(1, dtype=pto.i32) @@ -229,6 +255,14 @@ def main() -> None: assert conditional_lexical_text.count("pto.section.cube {") == 1 assert "scf.if" in conditional_lexical_text + sibling_conditional_text = lexical_section_sibling_conditional_rebinding_probe.compile().mlir_text() + assert sibling_conditional_text.count("pto.section.cube {") == 1 + assert sibling_conditional_text.count("pto.section.vector {") == 1 + assert sibling_conditional_text.count("scf.if") == 2 + with make_context() as context: + module = Module.parse(sibling_conditional_text, context) + module.operation.verify() + loop_carry_text = lexical_section_loop_carry_probe.compile().mlir_text() assert loop_carry_text.count("pto.section.cube {") == 1 assert "scf.for" in loop_carry_text From 8b937f5b57b1ed8c54792b9d21b6656ce4a85fb7 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 14:28:34 +0800 Subject: [PATCH 021/122] feat(ptodsl): add deprecated decorator --- ptodsl/ptodsl/_diagnostics.py | 50 +++++++++++++++++++++++++ ptodsl/ptodsl/pto.py | 2 +- ptodsl/tests/test_deprecated.py | 66 +++++++++++++++++++++++++++++++++ 3 files changed, 117 insertions(+), 1 deletion(-) create mode 100644 ptodsl/tests/test_deprecated.py diff --git a/ptodsl/ptodsl/_diagnostics.py b/ptodsl/ptodsl/_diagnostics.py index 8e8f7a62e9..08510ddf3c 100644 --- a/ptodsl/ptodsl/_diagnostics.py +++ b/ptodsl/ptodsl/_diagnostics.py @@ -9,6 +9,54 @@ from __future__ import annotations +import warnings +from functools import wraps +from typing import Callable, ParamSpec, TypeVar + + +P = ParamSpec("P") +R = TypeVar("R") + + +class PTODSLDeprecationWarning(UserWarning): + """Warning emitted when a deprecated PTODSL interface is called.""" + + +def deprecated(reason: str) -> Callable[[Callable[P, R]], Callable[P, R]]: + """Mark a PTODSL callable as deprecated and warn when it is called. + + ``reason`` should describe the replacement or migration path, for example:: + + @deprecated("use pto.vadd(vector, scalar, mask) instead") + def vadds(vector, scalar, mask): + ... + + The wrapper preserves the decorated callable's metadata and exposes a + ``__deprecated__`` marker for tooling. ``stacklevel=2`` points the warning + at the caller of the deprecated interface rather than at this wrapper. + """ + + if not isinstance(reason, str) or not reason.strip(): + raise TypeError("deprecated() requires a non-empty string reason") + + def decorate(function: Callable[P, R]) -> Callable[P, R]: + if not callable(function): + raise TypeError("deprecated() can only decorate a callable") + + @wraps(function) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + warnings.warn( + f"{function.__qualname__} is deprecated; {reason}", + PTODSLDeprecationWarning, + stacklevel=2, + ) + return function(*args, **kwargs) + + wrapper.__deprecated__ = reason + return wrapper + + return decorate + class PTODSLTracingMisuseError(TypeError): """Raised when authored Python misuses PTODSL runtime values during tracing.""" @@ -540,7 +588,9 @@ def unsupported_public_surface_error(name: str) -> AttributeError: __all__ = [ + "PTODSLDeprecationWarning", "PTODSLTracingMisuseError", + "deprecated", "explicit_mode_required_error", "explicit_mode_required_with_context_error", "host_tensor_metadata_error", diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index 5cb8b96168..1ba043f100 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -20,7 +20,7 @@ internally as ``_pto`` (``from ptoas.mlir.dialects import pto as _pto``). """ -from ._diagnostics import unsupported_public_surface_error +from ._diagnostics import deprecated, unsupported_public_surface_error # ── Types ───────────────────────────────────────────────────────────────────── from ._types import ( # noqa: F401 diff --git a/ptodsl/tests/test_deprecated.py b/ptodsl/tests/test_deprecated.py new file mode 100644 index 0000000000..b2817fd385 --- /dev/null +++ b/ptodsl/tests/test_deprecated.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software; you can redistribute it and/or modify it under the terms of +# the CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for the full text of the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR +# FITNESS FOR A PARTICULAR PURPOSE. See the License for the specific language governing permissions +# and limitations under the License. + +import inspect +import warnings + +from ptodsl import pto +from ptodsl._diagnostics import PTODSLDeprecationWarning, deprecated + + +@deprecated("use replacement() instead") +def old_function(value, *, scale=1): + """Return a scaled value.""" + return value * scale + + +def expect(condition: bool, message: str) -> None: + if not condition: + raise AssertionError(message) + + +def main() -> None: + expect(pto.deprecated is deprecated, "deprecated should be exported on pto") + expect(old_function.__name__ == "old_function", "decorator should preserve function metadata") + expect(old_function.__doc__ == "Return a scaled value.", "decorator should preserve the docstring") + expect( + str(inspect.signature(old_function)) == "(value, *, scale=1)", + "decorator should preserve the callable signature", + ) + expect( + old_function.__deprecated__ == "use replacement() instead", + "decorator should expose its migration reason", + ) + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + result = old_function(3, scale=2) + + expect(result == 6, "decorator should preserve the wrapped result") + expect(len(captured) == 1, "deprecated call should emit one warning") + warning = captured[0] + expect(warning.category is PTODSLDeprecationWarning, "warning should use the PTODSL category") + expect("old_function is deprecated" in str(warning.message), "warning should name the old API") + expect("use replacement() instead" in str(warning.message), "warning should include the migration reason") + expect(warning.filename == __file__, "warning should point at the caller") + + for invalid_reason in ("", 123): + try: + deprecated(invalid_reason) + except TypeError: + pass + else: + raise AssertionError("deprecated() should reject an invalid reason") + + print("ptodsl_deprecated: PASS") + + +if __name__ == "__main__": + main() From 8ed921aa2f74a95f6344ad7851c4d6cbc307a5b0 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 15:15:06 +0800 Subject: [PATCH 022/122] feat(vmi): unify vector scalar binary operations --- .../14-vmi-virtual-instruction-set.md | 37 +++-- ptodsl/examples/vci_subvl_group_launch.py | 2 +- ptodsl/examples/vci_vadds_share_launch.py | 2 +- ptodsl/ptodsl/_vmi_namespace.py | 90 ++++++++++-- ptodsl/tests/test_deprecated.py | 13 +- ptodsl/tests/test_jit_compile.py | 12 +- ptodsl/tests/test_vmi_binary_ops.py | 132 ++++++++++++++++++ ptodsl/tests/test_vmi_vshr_signedness.py | 4 +- 8 files changed, 252 insertions(+), 40 deletions(-) create mode 100644 ptodsl/tests/test_vmi_binary_ops.py diff --git a/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md b/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md index 5c89195475..32f111f539 100644 --- a/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md +++ b/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md @@ -570,15 +570,20 @@ operands. They form the arithmetic core of VMI SIMD kernels. #### `pto.vmi.vshl(lhs, rhs, mask=None, *, pmode=None) -> VRegType` #### `pto.vmi.vshr(lhs, rhs, mask=None, *, pmode=None) -> VRegType` -**Description**: Element-wise binary operation: `result[i] = lhs[i] rhs[i]` -for lanes where `mask[i]` is true (or all lanes when `mask` is omitted). +**Description**: These are element-wise binary operations. For `pto.vmi.vadd`, +when `rhs` is a VMI vector, `result[i] = lhs[i] + rhs[i]` and the VMI `vadd` +operation is emitted. When `rhs` is a scalar, the scalar is applied to every +lane and the VMI `vadds` operation is emitted. The other operations require a +VMI vector `rhs`. Operations are restricted to lanes where `mask[i]` is true +(or all lanes when `mask` is omitted and the selected form permits an omitted +mask). **Parameters**: | Parameter | Type | Description | |-----------|------|-------------| | `lhs` | `VRegType` | First operand vector | -| `rhs` | `VRegType` | Second operand vector | +| `rhs` | `VRegType` or `ScalarType` | Second vector operand or scalar addend | | `mask` | VMI mask or `None` | Optional predicate mask gating lane participation | | `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | @@ -596,7 +601,9 @@ out = pto.vmi.vmul(scale, data, full_mask) ``` **Constraints**: -- `lhs` and `rhs` must have compatible shapes and element types. +- For vector-vector form, `lhs` and `rhs` must have compatible shapes and + element types. +- For vector-scalar form, the scalar is coerced to the element type of `lhs`. - The result type is inferred from `lhs`. - For bitwise ops (`vand`, `vor`, `vxor`, `vshl`, `vshr`), integer element types are expected. Floating-point usage is rejected. @@ -646,12 +653,20 @@ inverted = pto.vmi.vnot(int_vec) Formal `pto.vmi` vector-scalar ops in VMI v0.1: -#### `pto.vmi.vadds(source, scalar, mask, *, pmode=None) -> VRegType` -#### `pto.vmi.vmuls(source, scalar, mask, *, pmode=None) -> VRegType` -#### `pto.vmi.vmaxs(source, scalar, mask, *, pmode=None) -> VRegType` -#### `pto.vmi.vmins(source, scalar, mask, *, pmode=None) -> VRegType` -#### `pto.vmi.vshls(source, scalar, mask, *, pmode=None) -> VRegType` -#### `pto.vmi.vshrs(source, scalar, mask, *, pmode=None) -> VRegType` +#### `pto.vmi.vadds(source, scalar, mask, *, pmode=None) -> VRegType` (deprecated) +#### `pto.vmi.vmuls(source, scalar, mask, *, pmode=None) -> VRegType` (deprecated) +#### `pto.vmi.vmaxs(source, scalar, mask, *, pmode=None) -> VRegType` (deprecated) +#### `pto.vmi.vmins(source, scalar, mask, *, pmode=None) -> VRegType` (deprecated) +#### `pto.vmi.vshls(source, scalar, mask, *, pmode=None) -> VRegType` (deprecated) +#### `pto.vmi.vshrs(source, scalar, mask, *, pmode=None) -> VRegType` (deprecated) + +These `*s` functions remain available as compatibility entry points and emit a +`PTODSLDeprecationWarning`. Use the matching unified entry point for new +PTODSL code, for example `pto.vmi.vmul(source, scalar, mask)` or +`pto.vmi.vshr(source, scalar, mask)`. The warning applies only to the Python +compatibility entry points; the underlying VMI `vadds`, `vmuls`, `vmaxs`, +`vmins`, `vshls`, and `vshrs` operations remain part of the instruction set +and are unchanged. The following are **PTODSL syntax sugar** — convenience wrappers provided by the PTODSL authoring layer. They have **no corresponding VMI instruction**; PTODSL lowers @@ -1452,7 +1467,7 @@ def vmi_elementwise( | Index / Broadcast | `vci`, `vbrc` | | Binary vector-vector | `vadd`, `vsub`, `vmul`, `vdiv`, `vmax`, `vmin`, `vand`, `vor`, `vxor`, `vshl`, `vshr` | | Unary vector | `vabs`, `vneg`, `vrelu`, `vexp`, `vln`, `vsqrt`, `vnot` | -| Vector-scalar | formal `pto.vmi`: `vadds`, `vmuls`, `vmaxs`, `vmins`, `vshls`, `vshrs`; DSL convenience: `vsubs`, `vands`, `vors`, `vxors` | +| Vector-scalar | `pto.vmi.vadd(vector, scalar, mask)` emits `vadds`; other formal `pto.vmi` helpers are `vmuls`, `vmaxs`, `vmins`, `vshls`, `vshrs`; DSL convenience: `vsubs`, `vands`, `vors`, `vxors` | | Compare / Select | `vcmp`, `vcmps`, `vsel`, `vselr` | | Reduction | `vcadd`, `vcmax`, `vcmin` | | Conversion | `vcvt`, `vinterpret_cast` | diff --git a/ptodsl/examples/vci_subvl_group_launch.py b/ptodsl/examples/vci_subvl_group_launch.py index 06e91e5685..05d2e49289 100644 --- a/ptodsl/examples/vci_subvl_group_launch.py +++ b/ptodsl/examples/vci_subvl_group_launch.py @@ -65,7 +65,7 @@ def kernel(out_ptr: pto.ptr(dtype, "gm")): ub = pto.castptr(pto.i64(0), pto.ptr(dtype, "ub")) mask = pto.vmi.create_mask(size, size=size) idx = pto.vmi.vci(dtype(0), size=size, group=group) - out_idx = pto.vmi.vadds(idx, dtype(ADD_SCALAR), mask) + out_idx = pto.vmi.vadd(idx, dtype(ADD_SCALAR), mask) pto.vmi.vstore(out_idx, ub, pto.const(0, dtype=pto.index)) pto.set_flag("V", "MTE3", event_id=0) pto.wait_flag("V", "MTE3", event_id=0) diff --git a/ptodsl/examples/vci_vadds_share_launch.py b/ptodsl/examples/vci_vadds_share_launch.py index 3263d592bb..7ada54a712 100644 --- a/ptodsl/examples/vci_vadds_share_launch.py +++ b/ptodsl/examples/vci_vadds_share_launch.py @@ -64,7 +64,7 @@ def kernel(out_ptr: pto.ptr(pto.i32, "gm")): mask = pto.vmi.create_mask(vl, size=vl) for _k in range(K_REMAT): idx = pto.vmi.vci(pto.i32(0), size=vl, group=num_groups) - out_idx = pto.vmi.vadds(idx, pto.i32(ADD_SCALAR), mask) + out_idx = pto.vmi.vadd(idx, pto.i32(ADD_SCALAR), mask) pto.vmi.vstore(out_idx, ub, pto.const(0, dtype=pto.index)) pto.set_flag("V", "MTE3", event_id=0) pto.wait_flag("V", "MTE3", event_id=0) diff --git a/ptodsl/ptodsl/_vmi_namespace.py b/ptodsl/ptodsl/_vmi_namespace.py index 68a9c7f807..09e46de1f2 100644 --- a/ptodsl/ptodsl/_vmi_namespace.py +++ b/ptodsl/ptodsl/_vmi_namespace.py @@ -25,6 +25,7 @@ ) from ._scalar_coercion import coerce_scalar_to_type +from ._diagnostics import deprecated from ._surface_values import _coerce_index_value, _try_get_constant_index, unwrap_surface_value, wrap_surface_value from ._types import ( VMI_LANE_COUNTS, @@ -565,6 +566,21 @@ def _emit_vec_scalar(op_name: str, source, scalar, mask, *, pmode=None, loc=None ) +def _emit_binary_or_vec_scalar( + binary_op_name: str, + vec_scalar_op_name: str, + lhs, + rhs, + mask=None, + **kw, +): + """Dispatch a VMI binary family from the second operand kind.""" + rhs_type = getattr(_raw(rhs), "type", None) + if rhs_type is not None and _is_vmi_vreg_type(rhs_type): + return _emit_binary(binary_op_name, lhs, rhs, mask, **kw) + return _emit_vec_scalar(vec_scalar_op_name, lhs, rhs, mask, **kw) + + def _emit_reduce( op_name: str, source, @@ -735,17 +751,43 @@ def vci(base, *, size, order=None, group=None, loc=None, ip=None): "vci", result_type, base, order=order, group=group, loc=loc, ip=ip ) - vadd = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vadd", lhs, rhs, mask, **kw)) + @staticmethod + def vadd(lhs, rhs, mask=None, **kw): + """Emit VMI vector addition, selecting vector or scalar form by type.""" + return _emit_binary_or_vec_scalar("vadd", "vadds", lhs, rhs, mask, **kw) + vsub = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vsub", lhs, rhs, mask, **kw)) - vmul = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vmul", lhs, rhs, mask, **kw)) + + @staticmethod + def vmul(lhs, rhs, mask=None, **kw): + """Emit VMI vector multiplication, selecting vector or scalar form by type.""" + return _emit_binary_or_vec_scalar("vmul", "vmuls", lhs, rhs, mask, **kw) + vdiv = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vdiv", lhs, rhs, mask, **kw)) - vmax = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vmax", lhs, rhs, mask, **kw)) - vmin = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vmin", lhs, rhs, mask, **kw)) + + @staticmethod + def vmax(lhs, rhs, mask=None, **kw): + """Emit VMI maximum, selecting vector or scalar form by type.""" + return _emit_binary_or_vec_scalar("vmax", "vmaxs", lhs, rhs, mask, **kw) + + @staticmethod + def vmin(lhs, rhs, mask=None, **kw): + """Emit VMI minimum, selecting vector or scalar form by type.""" + return _emit_binary_or_vec_scalar("vmin", "vmins", lhs, rhs, mask, **kw) + vand = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vand", lhs, rhs, mask, **kw)) vor = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vor", lhs, rhs, mask, **kw)) vxor = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vxor", lhs, rhs, mask, **kw)) - vshl = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vshl", lhs, rhs, mask, **kw)) - vshr = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vshr", lhs, rhs, mask, **kw)) + + @staticmethod + def vshl(lhs, rhs, mask=None, **kw): + """Emit VMI shift-left, selecting vector or scalar form by type.""" + return _emit_binary_or_vec_scalar("vshl", "vshls", lhs, rhs, mask, **kw) + + @staticmethod + def vshr(lhs, rhs, mask=None, **kw): + """Emit VMI shift-right, selecting vector or scalar form by type.""" + return _emit_binary_or_vec_scalar("vshr", "vshrs", lhs, rhs, mask, **kw) vabs = staticmethod(lambda source, mask=None, **kw: _emit_unary("vabs", source, mask, **kw)) vneg = staticmethod(lambda source, mask=None, **kw: _emit_unary("vneg", source, mask, **kw)) @@ -755,12 +797,36 @@ def vci(base, *, size, order=None, group=None, loc=None, ip=None): vsqrt = staticmethod(lambda source, mask=None, **kw: _emit_unary("vsqrt", source, mask, **kw)) vnot = staticmethod(lambda source, mask=None, **kw: _emit_unary("vnot", source, mask, **kw)) - vadds = staticmethod(lambda source, scalar, mask, **kw: _emit_vec_scalar("vadds", source, scalar, mask, **kw)) - vmuls = staticmethod(lambda source, scalar, mask, **kw: _emit_vec_scalar("vmuls", source, scalar, mask, **kw)) - vmaxs = staticmethod(lambda source, scalar, mask, **kw: _emit_vec_scalar("vmaxs", source, scalar, mask, **kw)) - vmins = staticmethod(lambda source, scalar, mask, **kw: _emit_vec_scalar("vmins", source, scalar, mask, **kw)) - vshls = staticmethod(lambda source, scalar, mask, **kw: _emit_vec_scalar("vshls", source, scalar, mask, **kw)) - vshrs = staticmethod(lambda source, scalar, mask, **kw: _emit_vec_scalar("vshrs", source, scalar, mask, **kw)) + @staticmethod + @deprecated("use pto.vmi.vadd(vector, scalar, mask) instead") + def vadds(source, scalar, mask, **kw): + """Deprecated VMI vector-scalar add compatibility entry point.""" + return _emit_vec_scalar("vadds", source, scalar, mask, **kw) + + @staticmethod + @deprecated("use pto.vmi.vmul(vector, scalar, mask) instead") + def vmuls(source, scalar, mask, **kw): + return _emit_vec_scalar("vmuls", source, scalar, mask, **kw) + + @staticmethod + @deprecated("use pto.vmi.vmax(vector, scalar, mask) instead") + def vmaxs(source, scalar, mask, **kw): + return _emit_vec_scalar("vmaxs", source, scalar, mask, **kw) + + @staticmethod + @deprecated("use pto.vmi.vmin(vector, scalar, mask) instead") + def vmins(source, scalar, mask, **kw): + return _emit_vec_scalar("vmins", source, scalar, mask, **kw) + + @staticmethod + @deprecated("use pto.vmi.vshl(vector, scalar, mask) instead") + def vshls(source, scalar, mask, **kw): + return _emit_vec_scalar("vshls", source, scalar, mask, **kw) + + @staticmethod + @deprecated("use pto.vmi.vshr(vector, scalar, mask) instead") + def vshrs(source, scalar, mask, **kw): + return _emit_vec_scalar("vshrs", source, scalar, mask, **kw) @staticmethod def vcmp(lhs, rhs, seed, cmp, *, pmode=None, loc=None, ip=None): diff --git a/ptodsl/tests/test_deprecated.py b/ptodsl/tests/test_deprecated.py index b2817fd385..223413d4ba 100644 --- a/ptodsl/tests/test_deprecated.py +++ b/ptodsl/tests/test_deprecated.py @@ -1,12 +1,11 @@ #!/usr/bin/env python3 # Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software; you can redistribute it and/or modify it under the terms of -# the CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for the full text of the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, -# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR -# FITNESS FOR A PARTICULAR PURPOSE. See the License for the specific language governing permissions -# and limitations under the License. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. import inspect import warnings diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 45dc7975f8..92c0ca4b37 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -2650,12 +2650,12 @@ def vmi_wrapper_dispatch_probe(): int_not = pto.vmi.vnot(int_lhs, mask) int_shl = pto.vmi.vshl(int_lhs, int_rhs, mask) int_shr = pto.vmi.vshr(int_lhs, int_rhs, mask) - scalar_added = pto.vmi.vadds(relu, 1.0, mask) - scalar_multiplied = pto.vmi.vmuls(relu, 2.0, mask) - scalar_maximum = pto.vmi.vmaxs(relu, 1.0, mask) - scalar_minimum = pto.vmi.vmins(relu, 1.0, mask) - scalar_shl = pto.vmi.vshls(int_lhs, pto.i32(1), mask) - scalar_shr = pto.vmi.vshrs(int_lhs, pto.i32(1), mask) + scalar_added = pto.vmi.vadd(relu, 1.0, mask) + scalar_multiplied = pto.vmi.vmul(relu, 2.0, mask) + scalar_maximum = pto.vmi.vmax(relu, 1.0, mask) + scalar_minimum = pto.vmi.vmin(relu, 1.0, mask) + scalar_shl = pto.vmi.vshl(int_lhs, pto.i32(1), mask) + scalar_shr = pto.vmi.vshr(int_lhs, pto.i32(1), mask) scaled = pto.vmi.vadd(scalar_multiplied, bias, mask) pred = pto.vmi.vcmp(scaled, lhs, mask, "ogt") scalar_pred = pto.vmi.vcmps(scaled, 0.0, mask, "ogt") diff --git a/ptodsl/tests/test_vmi_binary_ops.py b/ptodsl/tests/test_vmi_binary_ops.py new file mode 100644 index 0000000000..83ab8a0490 --- /dev/null +++ b/ptodsl/tests/test_vmi_binary_ops.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software; you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import warnings + +from ptodsl import pto +from ptodsl._diagnostics import PTODSLDeprecationWarning + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_binary_vector_vector_probe(): + lhs_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + rhs_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + mask = pto.vmi.create_mask(64, size=64) + lhs = pto.vmi.vload(lhs_tile.as_ptr(), 0, size=64) + rhs = pto.vmi.vload(rhs_tile.as_ptr(), 0, size=64) + _ = pto.vmi.vadd(lhs, rhs, mask) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_binary_add_vector_scalar_probe(): + source_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + mask = pto.vmi.create_mask(64, size=64) + source = pto.vmi.vload(source_tile.as_ptr(), 0, size=64) + _ = pto.vmi.vadd(source, 1.0, mask) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_binary_vector_scalar_probe(): + source_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + integer_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.i32) + mask = pto.vmi.create_mask(64, size=64) + source = pto.vmi.vload(source_tile.as_ptr(), 0, size=64) + integer_source = pto.vmi.vload(integer_tile.as_ptr(), 0, size=64) + _ = pto.vmi.vmul(source, 2.0, mask) + _ = pto.vmi.vmax(source, 1.0, mask) + _ = pto.vmi.vmin(source, 1.0, mask) + _ = pto.vmi.vshl(integer_source, 1, mask) + _ = pto.vmi.vshr(integer_source, 1, mask) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_deprecated_binary_scalar_probe(): + source_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + integer_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.i32) + mask = pto.vmi.create_mask(64, size=64) + source = pto.vmi.vload(source_tile.as_ptr(), 0, size=64) + integer_source = pto.vmi.vload(integer_tile.as_ptr(), 0, size=64) + _ = pto.vmi.vadds(source, 1.0, mask) + _ = pto.vmi.vmuls(source, 2.0, mask) + _ = pto.vmi.vmaxs(source, 1.0, mask) + _ = pto.vmi.vmins(source, 1.0, mask) + _ = pto.vmi.vshls(integer_source, 1, mask) + _ = pto.vmi.vshrs(integer_source, 1, mask) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_binary_add_compatibility_probe(): + source_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + mask = pto.vmi.create_mask(64, size=64) + source = pto.vmi.vload(source_tile.as_ptr(), 0, size=64) + _ = pto.vmi.vadds(source, 1.0, mask) + + +def expect(condition: bool, message: str) -> None: + if not condition: + raise AssertionError(message) + + +def main() -> None: + vector_vector_text = vmi_binary_vector_vector_probe.compile().mlir_text() + expect( + "pto.vmi.vadd" in vector_vector_text, + "vmi.vadd(vector, vector, mask) should emit pto.vmi.vadd", + ) + expect( + "pto.vmi.vadds" not in vector_vector_text, + "vector-vector vmi.vadd should not emit pto.vmi.vadds", + ) + + vector_scalar_text = vmi_binary_add_vector_scalar_probe.compile().mlir_text() + expect( + "pto.vmi.vadds" in vector_scalar_text, + "vmi.vadd(vector, scalar, mask) should emit pto.vmi.vadds", + ) + + vector_scalar_text = vmi_binary_vector_scalar_probe.compile().mlir_text() + for op_name in ("vmuls", "vmaxs", "vmins", "vshls", "vshrs"): + expect( + f"pto.vmi.{op_name}" in vector_scalar_text, + f"unified VMI scalar form should emit pto.vmi.{op_name}", + ) + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always", PTODSLDeprecationWarning) + compatibility_text = vmi_binary_add_compatibility_probe.compile().mlir_text() + + expect("pto.vmi.vadds" in compatibility_text, "vmi.vadds compatibility should preserve the VMI op") + deprecation_warnings = [ + warning for warning in captured if warning.category is PTODSLDeprecationWarning + ] + expect(len(deprecation_warnings) == 1, "vmi.vadds should emit one deprecation warning") + expect( + "pto.vmi.vadd(vector, scalar, mask)" in str(deprecation_warnings[0].message), + "vmi.vadds warning should name the replacement API", + ) + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always", PTODSLDeprecationWarning) + deprecated_text = vmi_deprecated_binary_scalar_probe.compile().mlir_text() + + for op_name in ("vadds", "vmuls", "vmaxs", "vmins", "vshls", "vshrs"): + expect( + f"pto.vmi.{op_name}" in deprecated_text, + f"deprecated VMI {op_name} should preserve the underlying operation", + ) + deprecation_warnings = [ + warning for warning in captured if warning.category is PTODSLDeprecationWarning + ] + expect(len(deprecation_warnings) == 6, "all deprecated VMI scalar compatibility APIs should warn") + + print("ptodsl_vmi_binary_ops: PASS") + + +if __name__ == "__main__": + main() diff --git a/ptodsl/tests/test_vmi_vshr_signedness.py b/ptodsl/tests/test_vmi_vshr_signedness.py index 19998c0e60..d755b45c0e 100644 --- a/ptodsl/tests/test_vmi_vshr_signedness.py +++ b/ptodsl/tests/test_vmi_vshr_signedness.py @@ -25,7 +25,7 @@ def vmi_vshr_signed_probe(): lhs = pto.vmi.vload(lhs_tile.as_ptr(), offset, size=128) rhs = pto.vmi.vload(rhs_tile.as_ptr(), offset, size=128) shifted = pto.vmi.vshr(lhs, rhs, mask) - shifted_scalar = pto.vmi.vshrs(lhs, pto.si32(3), mask) + shifted_scalar = pto.vmi.vshr(lhs, pto.si32(3), mask) _ = shifted _ = shifted_scalar @@ -40,7 +40,7 @@ def vmi_vshr_unsigned_probe(): lhs = pto.vmi.vload(lhs_tile.as_ptr(), offset, size=128) rhs = pto.vmi.vload(rhs_tile.as_ptr(), offset, size=128) shifted = pto.vmi.vshr(lhs, rhs, mask) - shifted_scalar = pto.vmi.vshrs(lhs, pto.ui32(3), mask) + shifted_scalar = pto.vmi.vshr(lhs, pto.ui32(3), mask) _ = shifted _ = shifted_scalar From 8f026cb99b6af4ecc646e2865abaa1eeb02d9180 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 15:21:39 +0800 Subject: [PATCH 023/122] fix(ptodsl): correct test license header --- ptodsl/tests/test_vmi_binary_ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ptodsl/tests/test_vmi_binary_ops.py b/ptodsl/tests/test_vmi_binary_ops.py index 83ab8a0490..b9b877170a 100644 --- a/ptodsl/tests/test_vmi_binary_ops.py +++ b/ptodsl/tests/test_vmi_binary_ops.py @@ -1,9 +1,9 @@ #!/usr/bin/env python3 # Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software; you can redistribute it and/or modify it under the terms and conditions of +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of # CANN Open Software License Agreement Version 2.0 (the "License"). # Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. From 18fe4c6efb69f12328de014e70c2dad85fb811f2 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 15:42:32 +0800 Subject: [PATCH 024/122] fix(ptodsl): preserve section entry values in conditional merges --- ptodsl/ptodsl/_ast_rewrite.py | 2 +- ptodsl/tests/test_section.py | 37 +++++++++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index c07e863e42..f7adebc04a 100644 --- a/ptodsl/ptodsl/_ast_rewrite.py +++ b/ptodsl/ptodsl/_ast_rewrite.py @@ -1211,7 +1211,7 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con result.extend( ast.Assign( targets=[_name(old_name, ast.Store())], - value=_name(name), + value=_name(self._section_entry_bindings.get(name, name)), ) for name, old_name in old_value_names.items() ) diff --git a/ptodsl/tests/test_section.py b/ptodsl/tests/test_section.py index 60c993cccb..655ca86b74 100644 --- a/ptodsl/tests/test_section.py +++ b/ptodsl/tests/test_section.py @@ -158,6 +158,35 @@ def lexical_section_sibling_conditional_rebinding_probe(): pto.wait_flag("MTE2", "S", event_id=n_tile) +@pto.jit(target="a5", mode="explicit") +def lexical_section_sibling_single_sided_conditional_rebinding_probe(): + one = pto.const(1, dtype=pto.i32) + m_tile = pto.const(0, dtype=pto.i32) + n_tile = pto.const(0, dtype=pto.i32) + with pto.section("cube"): + if pto.get_block_idx() < one: + m_tile = one + n_tile = one + m_tile_1 = pto.const(0, dtype=pto.i32) + n_tile_1 = pto.const(0, dtype=pto.i32) + if pto.get_block_idx() < one: + m_tile_1 = m_tile + n_tile_1 = n_tile + pto.wait_flag("S", "MTE2", event_id=m_tile_1) + pto.wait_flag("S", "MTE2", event_id=n_tile_1) + with pto.section("vector"): + if pto.get_block_idx() < one: + m_tile = one + n_tile = one + m_tile_2 = pto.const(0, dtype=pto.i32) + n_tile_2 = pto.const(0, dtype=pto.i32) + if pto.get_block_idx() < one: + m_tile_2 = m_tile + n_tile_2 = n_tile + pto.wait_flag("MTE2", "S", event_id=m_tile_2) + pto.wait_flag("MTE2", "S", event_id=n_tile_2) + + @pto.jit(target="a5", mode="explicit") def lexical_section_loop_carry_probe(): one = pto.const(1, dtype=pto.i32) @@ -263,6 +292,14 @@ def main() -> None: module = Module.parse(sibling_conditional_text, context) module.operation.verify() + sibling_single_sided_text = lexical_section_sibling_single_sided_conditional_rebinding_probe.compile().mlir_text() + assert sibling_single_sided_text.count("pto.section.cube {") == 1 + assert sibling_single_sided_text.count("pto.section.vector {") == 1 + assert sibling_single_sided_text.count("scf.if") == 4 + with make_context() as context: + module = Module.parse(sibling_single_sided_text, context) + module.operation.verify() + loop_carry_text = lexical_section_loop_carry_probe.compile().mlir_text() assert loop_carry_text.count("pto.section.cube {") == 1 assert "scf.for" in loop_carry_text From ef95f386470ce5492b796613b031e0d2abedeea0 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Thu, 6 Aug 2026 16:34:32 +0800 Subject: [PATCH 025/122] fix: preserve A5 tinsert fp tile semantics --- ...todsl-tilelib-template-selection-design.md | 1 + lib/PTO/Transforms/ExpandTileOp.cpp | 8 ++ .../Transforms/InsertTemplateAttributes.cpp | 8 ++ lib/TileOps/a5/tinsert.py | 87 ++++++++++++++----- .../tilelib-template-authoring.md | 2 + ...nd_tile_op_tilelang_tinsert_fp_acc2mat.pto | 83 +++++++++++++----- 6 files changed, 146 insertions(+), 43 deletions(-) diff --git a/docs/designs/ptodsl-tilelib-template-selection-design.md b/docs/designs/ptodsl-tilelib-template-selection-design.md index 16e8b7f503..fead7370a5 100644 --- a/docs/designs/ptodsl-tilelib-template-selection-design.md +++ b/docs/designs/ptodsl-tilelib-template-selection-design.md @@ -129,6 +129,7 @@ context attrs. Current examples include: | `cmp_mode` | `tcmp`, `tcmps` | | `mask_pattern` | gather-side paths | | `precisionType` | high-precision math families | +| `acc_to_vec_mode`, `relu_pre_mode` | `tinsert` accumulator writeback paths | When a new TileLangDSL version depends on an op attribute, the PTODSL migration should first decide whether the attribute is a real context attr. If it changes diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 1643e529f3..28607161ba 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -505,6 +505,14 @@ static LogicalResult appendOpContextAttrs( stringifyCmpMode(cmpModeAttr.getValue()).str()); } } + if (auto tinsert = dyn_cast(op)) { + if (auto modeAttr = tinsert.getAccToVecModeAttr()) { + attrs.emplace_back("acc_to_vec_mode", + stringifyAccToVecMode(modeAttr.getValue()).str()); + } + attrs.emplace_back("relu_pre_mode", + stringifyReluPreMode(tinsert.getReluPreMode()).str()); + } if (auto tgather = dyn_cast(op)) { if (auto maskPatternAttr = tgather.getMaskPatternAttr()) { attrs.emplace_back( diff --git a/lib/PTO/Transforms/InsertTemplateAttributes.cpp b/lib/PTO/Transforms/InsertTemplateAttributes.cpp index f1069d4f5b..b31a23b59f 100644 --- a/lib/PTO/Transforms/InsertTemplateAttributes.cpp +++ b/lib/PTO/Transforms/InsertTemplateAttributes.cpp @@ -502,6 +502,14 @@ static void appendOpContextAttrs( attrs.emplace_back("cmp_mode", stringifyCmpMode(cmpModeAttr.getValue()).str()); } + if (auto tinsert = dyn_cast(op)) { + if (auto modeAttr = tinsert.getAccToVecModeAttr()) { + attrs.emplace_back("acc_to_vec_mode", + stringifyAccToVecMode(modeAttr.getValue()).str()); + } + attrs.emplace_back("relu_pre_mode", + stringifyReluPreMode(tinsert.getReluPreMode()).str()); + } if (auto tmrgsort = dyn_cast(op)) attrs.emplace_back("exhausted", tmrgsort.getExhausted() ? "1" : "0"); diff --git a/lib/TileOps/a5/tinsert.py b/lib/TileOps/a5/tinsert.py index c0fab9d7cd..82eb14f7a5 100644 --- a/lib/TileOps/a5/tinsert.py +++ b/lib/TileOps/a5/tinsert.py @@ -80,24 +80,43 @@ def _vec_to_vec_nd_scalar(src_memory_space, dst_memory_space, src_config, dst_co _DTYPES = [(dtype, "i32", "i32", dtype) for dtype in NUMERIC_DTYPES] -_FP_DTYPES = ( - ("f32", "i32", "i32", "i8", "f32"), - ("f32", "i32", "i32", "si8", "f32"), - ("f32", "i32", "i32", "ui8", "f32"), - ("f32", "i32", "i32", "f16", "f32"), - ("f32", "i32", "i32", "bf16", "f32"), - ("f32", "i32", "i32", "f32", "f32"), - ("i32", "i32", "i32", "i8", "f32"), - ("si32", "i32", "i32", "i8", "f32"), - ("i32", "i32", "i32", "si8", "f32"), - ("si32", "i32", "i32", "si8", "f32"), - ("i32", "i32", "i32", "ui8", "f32"), - ("si32", "i32", "i32", "ui8", "f32"), - ("i32", "i32", "i32", "f16", "f32"), - ("si32", "i32", "i32", "f16", "f32"), - ("i32", "i32", "i32", "bf16", "f32"), - ("si32", "i32", "i32", "bf16", "f32"), +_FP_DATA_DTYPES = ( + ("f32", "i32", "i32", "i8"), + ("f32", "i32", "i32", "si8"), + ("f32", "i32", "i32", "ui8"), + ("f32", "i32", "i32", "f8e4m3"), + ("f32", "i32", "i32", "hif8"), + ("f32", "i32", "i32", "f16"), + ("f32", "i32", "i32", "bf16"), + ("f32", "i32", "i32", "f32"), + ("i32", "i32", "i32", "i8"), + ("si32", "i32", "i32", "i8"), + ("i32", "i32", "i32", "si8"), + ("si32", "i32", "i32", "si8"), + ("i32", "i32", "i32", "ui8"), + ("si32", "i32", "i32", "ui8"), + ("i32", "i32", "i32", "f16"), + ("si32", "i32", "i32", "f16"), + ("i32", "i32", "i32", "bf16"), + ("si32", "i32", "i32", "bf16"), ) +_FP_SCALING_DTYPES = ("f16", "bf16", "f32") +_FP_DTYPES = tuple( + (*signature, fp_dtype) + for signature in _FP_DATA_DTYPES + for fp_dtype in _FP_SCALING_DTYPES +) +_FP_NZ_DTYPES = tuple( + signature for signature in _FP_DTYPES if signature[3] == "f32" +) + + +def _canonical_dtype_name(dtype): + name = str(dtype) + return { + "f8E4M3FN": "f8e4m3", + "!pto.hif8": "hif8", + }.get(name, name) def _tinsert_fp_quant_mode(src_dtype, dst_dtype): @@ -105,6 +124,8 @@ def _tinsert_fp_quant_mode(src_dtype, dst_dtype): ("f32", "i8"): "qf322b8_pre_vec", ("f32", "si8"): "qf322b8_pre_vec", ("f32", "ui8"): "qf322b8_pre_vec", + ("f32", "f8e4m3"): "qf322fp8_pre_vec", + ("f32", "hif8"): "qf322hif8_pre_vec", ("f32", "f16"): "qf322f16_pre_vec", ("f32", "bf16"): "qf322bf16_pre_vec", ("f32", "f32"): "qf322f32_pre_vec", @@ -119,7 +140,21 @@ def _tinsert_fp_quant_mode(src_dtype, dst_dtype): ("i32", "bf16"): "qs322bf16_pre_vec", ("si32", "bf16"): "qs322bf16_pre_vec", } - return modes[(str(src_dtype), str(dst_dtype))] + return modes[ + (_canonical_dtype_name(src_dtype), _canonical_dtype_name(dst_dtype)) + ] + + +def _acc_src_ptr(src): + src_ptr = src.as_ptr() + if str(src.dtype) == "si32": + src_ptr = pto.castptr(src_ptr, pto.ptr(pto.i32, "acc")) + return src_ptr + + +def _tinsert_pre_relu(): + mode = pto.get_op_attr("relu_pre_mode", "no_relu") + return None if mode == "no_relu" else (mode, None, None) def _acc_to_mat_fp( @@ -252,6 +287,7 @@ def template_tinsert_acc_to_mat_basic( n_size, src.shape[0] * pto.bytewidth(src.dtype), dst.shape[0] * c0_size * elem_bytes, + pre_relu=_tinsert_pre_relu(), ) @@ -286,13 +322,14 @@ def template_tinsert_fp_acc_to_mat( dst_offset = dst.shape[0] * c0_size * col_block + index_row * c0_size + col_mod pto.mte_l0c_l1( - src.as_ptr(), + _acc_src_ptr(src), pto.addptr(dst.as_ptr(), dst_offset), valid_rows, n_size, src.shape[0] * pto.bytewidth(src.dtype), dst.shape[0] * c0_size * elem_bytes, pre_quant=(fp.as_ptr(), _tinsert_fp_quant_mode(src.dtype, dst.dtype)), + pre_relu=_tinsert_pre_relu(), ) @@ -340,6 +377,7 @@ def template_tinsert_acc_to_vec_nd_basic( dst.shape[1], dst_mode, pre_quant=(pto.f16(1.0), "f32_f16"), + pre_relu=_tinsert_pre_relu(), **kwargs, ) elif str(src.dtype) == "f32" and str(dst.dtype) == "bf16": @@ -352,6 +390,7 @@ def template_tinsert_acc_to_vec_nd_basic( dst.shape[1], dst_mode, pre_quant=(pto.bf16(1.0), "f32_bf16"), + pre_relu=_tinsert_pre_relu(), **kwargs, ) else: @@ -363,6 +402,7 @@ def template_tinsert_acc_to_vec_nd_basic( (valid_rows + 15) // 16 * 16, dst.shape[1], dst_mode, + pre_relu=_tinsert_pre_relu(), **kwargs, ) @@ -398,7 +438,7 @@ def template_tinsert_fp_acc_to_vec_nd( kwargs["layout"] = "nz2nd" pto.mte_l0c_ub( - src.as_ptr(), + _acc_src_ptr(src), dst_ptr, valid_rows, valid_cols, @@ -406,6 +446,7 @@ def template_tinsert_fp_acc_to_vec_nd( dst.shape[1], dst_mode, pre_quant=(fp.as_ptr(), _tinsert_fp_quant_mode(src.dtype, dst.dtype)), + pre_relu=_tinsert_pre_relu(), **kwargs, ) @@ -456,6 +497,7 @@ def template_tinsert_acc_to_vec_nz_basic( (valid_rows + 15) // 16 * 16 * pto.bytewidth(src.dtype), dst.shape[0] * c0_size * elem_bytes, dst_mode, + pre_relu=_tinsert_pre_relu(), **kwargs, ) @@ -464,7 +506,7 @@ def template_tinsert_acc_to_vec_nz_basic( op="pto.tinsert", target="a5", name="template_tinsert_fp_acc_to_vec_nz", - dtypes=_FP_DTYPES, + dtypes=_FP_NZ_DTYPES, iteration_axis="none", op_engine="other", op_class="movement", @@ -501,7 +543,7 @@ def template_tinsert_fp_acc_to_vec_nz( kwargs["layout"] = ("nz2nz", 0) pto.mte_l0c_ub( - src.as_ptr(), + _acc_src_ptr(src), pto.addptr(dst.as_ptr(), dst_offset), valid_rows, valid_cols, @@ -509,6 +551,7 @@ def template_tinsert_fp_acc_to_vec_nz( dst.shape[0] * c0_size * elem_bytes, dst_mode, pre_quant=(fp.as_ptr(), _tinsert_fp_quant_mode(src.dtype, dst.dtype)), + pre_relu=_tinsert_pre_relu(), **kwargs, ) diff --git a/ptodsl/docs/developer_guide/tilelib-template-authoring.md b/ptodsl/docs/developer_guide/tilelib-template-authoring.md index df05cefd6c..03e6a25ec9 100644 --- a/ptodsl/docs/developer_guide/tilelib-template-authoring.md +++ b/ptodsl/docs/developer_guide/tilelib-template-authoring.md @@ -110,6 +110,8 @@ forwards selected attributes as context attrs: - `cmp_mode` - `mask_pattern` - `precisionType` +- `acc_to_vec_mode` +- `relu_pre_mode` If a template needs a new op attribute, update the C++ context-attr forwarding before relying on it in Python. A template that silently assumes a default when diff --git a/test/lit/vpto/expand_tile_op_tilelang_tinsert_fp_acc2mat.pto b/test/lit/vpto/expand_tile_op_tilelang_tinsert_fp_acc2mat.pto index 02b828890c..ea9f1b2dc4 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tinsert_fp_acc2mat.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tinsert_fp_acc2mat.pto @@ -6,43 +6,84 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// The unified five-operand pto.tinsert fp form must be discoverable for every -// A5 Acc destination layout, not only by direct EmitC lowering. -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s +// The unified five-operand pto.tinsert fp form must preserve its attrs and +// accept every scaling dtype across the supported A5 Acc writeback forms. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s --check-prefix=FINAL -// CHECK-LABEL: func.func @TINSERT_FP_ACC_TO_MAT() -// CHECK-NOT: pto.tinsert -// CHECK: pto.set_fpc -// CHECK: pto.copy_matrix_cc_to_cbuf +// SELECT-DAG: pre_relu(mode = normal_relu) +// SELECT-DAG: pre_quant({{.*}}, mode = qf322fp8_pre_vec) +// SELECT-DAG: pre_quant({{.*}}, mode = qf322hif8_pre_vec) +// SELECT-DAG: pre_quant({{.*}}, mode = qs322bf16_pre_vec) +// SELECT-DAG: %[[VEC1:.*]] = arith.constant 1 : i64 +// SELECT-DAG: pto.mte_l0c_ub {{.*}} dst_mode(%[[VEC1]]) +// SELECT-DAG: pto.castptr {{.*}} : !pto.ptr -> !pto.ptr -// CHECK-LABEL: func.func @TINSERT_FP_ACC_TO_VEC_ND() -// CHECK-NOT: pto.tinsert -// CHECK: pto.set_fpc -// CHECK: pto.copy_matrix_cc_to_ub +// FINAL-LABEL: func.func @TINSERT_FP_ACC_TO_MAT_RELU_F16_SCALE() +// FINAL-NOT: pto.tinsert +// FINAL: pto.set_fpc +// FINAL: pto.copy_matrix_cc_to_cbuf -// CHECK-LABEL: func.func @TINSERT_FP_ACC_TO_VEC_NZ() -// CHECK-NOT: pto.tinsert -// CHECK: pto.set_fpc -// CHECK: pto.copy_matrix_cc_to_ub +// FINAL-LABEL: func.func @TINSERT_FP_ACC_TO_VEC_ND_FP8_BF16_SCALE() +// FINAL-NOT: pto.tinsert +// FINAL: pto.set_fpc +// FINAL: pto.copy_matrix_cc_to_ub + +// FINAL-LABEL: func.func @TINSERT_FP_ACC_TO_MAT_HIF8_F32_SCALE() +// FINAL-NOT: pto.tinsert +// FINAL: pto.set_fpc +// FINAL: pto.copy_matrix_cc_to_cbuf + +// FINAL-LABEL: func.func @TINSERT_FP_ACC_TO_VEC_ND_SI32() +// FINAL-NOT: pto.tinsert +// FINAL: pto.set_fpc +// FINAL: pto.copy_matrix_cc_to_ub + +// FINAL-LABEL: func.func @TINSERT_FP_ACC_TO_VEC_NZ() +// FINAL-NOT: pto.tinsert +// FINAL: pto.set_fpc +// FINAL: pto.copy_matrix_cc_to_ub module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { - func.func @TINSERT_FP_ACC_TO_MAT() attributes {pto.aicore} { + func.func @TINSERT_FP_ACC_TO_MAT_RELU_F16_SCALE() attributes {pto.aicore} { %c0 = arith.constant 0 : index %src = pto.alloc_tile : !pto.tile_buf - %fp = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) outs(%dst : !pto.tile_buf) + {reluPreMode = #pto} + return + } + + func.func @TINSERT_FP_ACC_TO_VEC_ND_FP8_BF16_SCALE() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {accToVecMode = #pto.acc_to_vec_mode} return } - func.func @TINSERT_FP_ACC_TO_VEC_ND() attributes {pto.aicore} { + func.func @TINSERT_FP_ACC_TO_MAT_HIF8_F32_SCALE() attributes {pto.aicore} { %c0 = arith.constant 0 : index %src = pto.alloc_tile : !pto.tile_buf %fp = pto.alloc_tile : !pto.tile_buf - %dst = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) - outs(%dst : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @TINSERT_FP_ACC_TO_VEC_ND_SI32() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index fp %fp : !pto.tile_buf) + outs(%dst : !pto.tile_buf) return } From c597a3e6d92218524f9609b4d716773d7aed86d7 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 16:59:09 +0800 Subject: [PATCH 026/122] feat(vmi): accept commutative scalar vector operands --- .../14-vmi-virtual-instruction-set.md | 6 ++++-- ptodsl/ptodsl/_vmi_namespace.py | 15 ++++++++++----- ptodsl/tests/test_vmi_binary_ops.py | 18 ++++++++++++++++++ 3 files changed, 32 insertions(+), 7 deletions(-) diff --git a/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md b/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md index 32f111f539..5f56fb2593 100644 --- a/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md +++ b/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md @@ -573,8 +573,10 @@ operands. They form the arithmetic core of VMI SIMD kernels. **Description**: These are element-wise binary operations. For `pto.vmi.vadd`, when `rhs` is a VMI vector, `result[i] = lhs[i] + rhs[i]` and the VMI `vadd` operation is emitted. When `rhs` is a scalar, the scalar is applied to every -lane and the VMI `vadds` operation is emitted. The other operations require a -VMI vector `rhs`. Operations are restricted to lanes where `mask[i]` is true +lane and the VMI `vadds` operation is emitted. For commutative operations +(`vadd`, `vmul`, `vmax`, and `vmin`), a scalar `lhs` with a vector `rhs` is also +accepted and normalized to the corresponding vector-scalar operation. The +other operations require a VMI vector `rhs`. Operations are restricted to lanes where `mask[i]` is true (or all lanes when `mask` is omitted and the selected form permits an omitted mask). diff --git a/ptodsl/ptodsl/_vmi_namespace.py b/ptodsl/ptodsl/_vmi_namespace.py index 09e46de1f2..01093ca10c 100644 --- a/ptodsl/ptodsl/_vmi_namespace.py +++ b/ptodsl/ptodsl/_vmi_namespace.py @@ -572,11 +572,16 @@ def _emit_binary_or_vec_scalar( lhs, rhs, mask=None, + *, + commutative=False, **kw, ): - """Dispatch a VMI binary family from the second operand kind.""" + """Dispatch a VMI binary family from the operand kinds.""" + lhs_type = getattr(_raw(lhs), "type", None) rhs_type = getattr(_raw(rhs), "type", None) if rhs_type is not None and _is_vmi_vreg_type(rhs_type): + if commutative and (lhs_type is None or not _is_vmi_vreg_type(lhs_type)): + return _emit_vec_scalar(vec_scalar_op_name, rhs, lhs, mask, **kw) return _emit_binary(binary_op_name, lhs, rhs, mask, **kw) return _emit_vec_scalar(vec_scalar_op_name, lhs, rhs, mask, **kw) @@ -754,26 +759,26 @@ def vci(base, *, size, order=None, group=None, loc=None, ip=None): @staticmethod def vadd(lhs, rhs, mask=None, **kw): """Emit VMI vector addition, selecting vector or scalar form by type.""" - return _emit_binary_or_vec_scalar("vadd", "vadds", lhs, rhs, mask, **kw) + return _emit_binary_or_vec_scalar("vadd", "vadds", lhs, rhs, mask, commutative=True, **kw) vsub = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vsub", lhs, rhs, mask, **kw)) @staticmethod def vmul(lhs, rhs, mask=None, **kw): """Emit VMI vector multiplication, selecting vector or scalar form by type.""" - return _emit_binary_or_vec_scalar("vmul", "vmuls", lhs, rhs, mask, **kw) + return _emit_binary_or_vec_scalar("vmul", "vmuls", lhs, rhs, mask, commutative=True, **kw) vdiv = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vdiv", lhs, rhs, mask, **kw)) @staticmethod def vmax(lhs, rhs, mask=None, **kw): """Emit VMI maximum, selecting vector or scalar form by type.""" - return _emit_binary_or_vec_scalar("vmax", "vmaxs", lhs, rhs, mask, **kw) + return _emit_binary_or_vec_scalar("vmax", "vmaxs", lhs, rhs, mask, commutative=True, **kw) @staticmethod def vmin(lhs, rhs, mask=None, **kw): """Emit VMI minimum, selecting vector or scalar form by type.""" - return _emit_binary_or_vec_scalar("vmin", "vmins", lhs, rhs, mask, **kw) + return _emit_binary_or_vec_scalar("vmin", "vmins", lhs, rhs, mask, commutative=True, **kw) vand = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vand", lhs, rhs, mask, **kw)) vor = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vor", lhs, rhs, mask, **kw)) diff --git a/ptodsl/tests/test_vmi_binary_ops.py b/ptodsl/tests/test_vmi_binary_ops.py index b9b877170a..47ccdf889b 100644 --- a/ptodsl/tests/test_vmi_binary_ops.py +++ b/ptodsl/tests/test_vmi_binary_ops.py @@ -45,6 +45,17 @@ def vmi_binary_vector_scalar_probe(): _ = pto.vmi.vshr(integer_source, 1, mask) +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_binary_scalar_vector_probe(): + source_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + mask = pto.vmi.create_mask(64, size=64) + source = pto.vmi.vload(source_tile.as_ptr(), 0, size=64) + _ = pto.vmi.vadd(1.0, source, mask) + _ = pto.vmi.vmul(2.0, source, mask) + _ = pto.vmi.vmax(1.0, source, mask) + _ = pto.vmi.vmin(1.0, source, mask) + + @pto.jit(target="a5", backend="vpto", mode="explicit") def vmi_deprecated_binary_scalar_probe(): source_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) @@ -97,6 +108,13 @@ def main() -> None: f"unified VMI scalar form should emit pto.vmi.{op_name}", ) + scalar_vector_text = vmi_binary_scalar_vector_probe.compile().mlir_text() + for op_name in ("vadds", "vmuls", "vmaxs", "vmins"): + expect( + f"pto.vmi.{op_name}" in scalar_vector_text, + f"commutative VMI scalar-vector form should emit pto.vmi.{op_name}", + ) + with warnings.catch_warnings(record=True) as captured: warnings.simplefilter("always", PTODSLDeprecationWarning) compatibility_text = vmi_binary_add_compatibility_probe.compile().mlir_text() From 799a55aed8e49348781160f4cc8a4f5292ad8063 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 20:27:18 +0800 Subject: [PATCH 027/122] fix(ptodsl): diagnose uninitialized section merge values --- ptodsl/ptodsl/_ast_rewrite.py | 31 ++++++++++++++++---- ptodsl/tests/test_section.py | 55 +++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 5 deletions(-) diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index f7adebc04a..787ee0d83d 100644 --- a/ptodsl/ptodsl/_ast_rewrite.py +++ b/ptodsl/ptodsl/_ast_rewrite.py @@ -59,6 +59,7 @@ def rewrite_jit_function(fn, *, static_bindings=None, rewrite_control_flow=True) rewriter = _ControlFlowRewriter( static_env, section_entry_bindings=section_rewriter.section_entry_bindings, + section_uninitialized_aliases=section_rewriter.section_uninitialized_aliases, ) function_def.body = rewriter.rewrite_block(function_def.body, live_after=set()) tree = ast.Module(body=[function_def], type_ignores=[]) @@ -101,6 +102,7 @@ def __init__(self): self._known_bindings = set() self._section_outer_bindings = None self.section_entry_bindings = {} + self.section_uninitialized_aliases = set() @staticmethod def _is_section_with(node): @@ -118,7 +120,9 @@ def _target_names(self, target): def _activate_targets(self, targets): for name in targets & self._local_names: - alias = self._env.setdefault(name, self._fresh_alias(name)) + if name not in self._env: + self._env[name] = self._fresh_alias(name) + alias = self._env[name] if self._section_outer_bindings is not None and name in self._section_outer_bindings: self.section_entry_bindings.setdefault(alias, name) @@ -201,13 +205,21 @@ def visit_For(self, node): def visit_If(self, node): node.test = self.visit(node.test) # Both branches of a runtime conditional share one authored binding. - # Reserve its section-local alias before visiting either branch so the - # branch merge does not treat the second branch as a new binding. + # Any future env-forking visitor must apply the same invariant: reserve + # common targets before visiting either branch. For section-local + # bindings this prevents the branch merge from creating two aliases. common_targets = _name_info(node.body).stores & _name_info(node.orelse).stores self._activate_targets(common_targets) entry_env = dict(self._env) node.body, body_env = self._visit_block(node.body, entry_env) node.orelse, else_env = self._visit_block(node.orelse, entry_env) + entry_aliases = set(entry_env.values()) + branch_only_aliases = set(body_env.values()) ^ set(else_env.values()) + self.section_uninitialized_aliases.update( + alias + for alias in branch_only_aliases - entry_aliases + if alias not in self.section_entry_bindings + ) self._env.update(body_env) self._env.update(else_env) return node @@ -947,9 +959,10 @@ def visit_Subscript(self, node): class _ControlFlowRewriter: - def __init__(self, static_env=None, *, section_entry_bindings=None): + def __init__(self, static_env=None, *, section_entry_bindings=None, section_uninitialized_aliases=None): self._static_env = dict(static_env or {}) self._section_entry_bindings = dict(section_entry_bindings or {}) + self._section_uninitialized_aliases = set(section_uninitialized_aliases or ()) self._counter = 0 def _fresh(self, prefix: str) -> str: @@ -957,6 +970,14 @@ def _fresh(self, prefix: str) -> str: self._counter += 1 return value + def _section_entry_value(self, name): + if name in self._section_uninitialized_aliases: + raise PTODSLAstRewriteError( + "ast_rewrite=True runtime if reads a section-local value before it is initialized; " + f"initialize {name!r} before the conditional" + ) + return _name(self._section_entry_bindings.get(name, name)) + def rewrite_block(self, stmts, *, live_after, live_after_slots=None, allow_loop_control=False, static_iters=None): rewritten_reversed = [] live = set(live_after) @@ -1211,7 +1232,7 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con result.extend( ast.Assign( targets=[_name(old_name, ast.Store())], - value=_name(self._section_entry_bindings.get(name, name)), + value=self._section_entry_value(name), ) for name, old_name in old_value_names.items() ) diff --git a/ptodsl/tests/test_section.py b/ptodsl/tests/test_section.py index 655ca86b74..5e7b132010 100644 --- a/ptodsl/tests/test_section.py +++ b/ptodsl/tests/test_section.py @@ -10,6 +10,7 @@ """Focused tracing coverage for explicit physical section hints.""" from ptodsl import pto +from ptodsl._ast_rewrite import PTODSLAstRewriteError from ptodsl._context import make_context from ptodsl._tracing.active import current_session from ptoas.mlir.ir import Module @@ -120,6 +121,17 @@ def lexical_section_rebinding_probe(): pto.wait_flag("MTE2", "S", event_id=event_id) +@pto.jit(target="a5") +def lexical_non_section_conditional_rebinding_probe(): + one = pto.const(1, dtype=pto.i32) + value = pto.const(0, dtype=pto.i32) + if pto.get_block_idx() < one: + value = one + else: + value = pto.const(2, dtype=pto.i32) + pto.wait_flag("S", "MTE2", event_id=value) + + @pto.jit(target="a5", mode="explicit") def lexical_section_conditional_rebinding_probe(): one = pto.const(1, dtype=pto.i32) @@ -187,6 +199,30 @@ def lexical_section_sibling_single_sided_conditional_rebinding_probe(): pto.wait_flag("MTE2", "S", event_id=n_tile_2) +@pto.jit(target="a5", mode="explicit") +def lexical_section_uninitialized_conditional_probe(): + one = pto.const(1, dtype=pto.i32) + with pto.section("cube"): + if pto.get_block_idx() < one: + value = one + pto.wait_flag("S", "MTE2", event_id=value) + + +@pto.jit(target="a5", mode="explicit") +def lexical_section_nested_conditional_rebinding_probe(): + one = pto.const(1, dtype=pto.i32) + value = pto.const(0, dtype=pto.i32) + with pto.section("cube"): + if pto.get_block_idx() < one: + if pto.get_block_idx() < one: + value = one + else: + value = pto.const(2, dtype=pto.i32) + else: + value = pto.const(3, dtype=pto.i32) + pto.wait_flag("S", "MTE2", event_id=value) + + @pto.jit(target="a5", mode="explicit") def lexical_section_loop_carry_probe(): one = pto.const(1, dtype=pto.i32) @@ -280,6 +316,12 @@ def main() -> None: assert lexical_text.count("pto.section.cube {") == 1 assert lexical_text.count("pto.section.vector {") == 1 + non_section_lexical_text = lexical_non_section_conditional_rebinding_probe.compile().mlir_text() + assert non_section_lexical_text.count("scf.if") == 1 + with make_context() as context: + module = Module.parse(non_section_lexical_text, context) + module.operation.verify() + conditional_lexical_text = lexical_section_conditional_rebinding_probe.compile().mlir_text() assert conditional_lexical_text.count("pto.section.cube {") == 1 assert "scf.if" in conditional_lexical_text @@ -300,6 +342,19 @@ def main() -> None: module = Module.parse(sibling_single_sided_text, context) module.operation.verify() + nested_conditional_text = lexical_section_nested_conditional_rebinding_probe.compile().mlir_text() + assert nested_conditional_text.count("pto.section.cube {") == 1 + assert nested_conditional_text.count("scf.if") == 2 + with make_context() as context: + module = Module.parse(nested_conditional_text, context) + module.operation.verify() + + _expect_raises( + PTODSLAstRewriteError, + lambda: lexical_section_uninitialized_conditional_probe.compile(), + "reads a section-local value before it is initialized", + ) + loop_carry_text = lexical_section_loop_carry_probe.compile().mlir_text() assert loop_carry_text.count("pto.section.cube {") == 1 assert "scf.for" in loop_carry_text From 749fc25cf1f8421b32186726637778a8db1c3d92 Mon Sep 17 00:00:00 2001 From: mouliangyu <21963576+mouliangyu@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:06:57 +0800 Subject: [PATCH 028/122] fix(vpto): keep block queries outside inferred vecscope (#1148) --- lib/PTO/Transforms/PTOInferVPTOVecScope.cpp | 6 +- ...to_vecscope_infer_block_query_boundary.pto | 72 +++++++++++++++++++ 2 files changed, 77 insertions(+), 1 deletion(-) create mode 100644 test/lit/vpto/auto_vecscope_infer_block_query_boundary.pto diff --git a/lib/PTO/Transforms/PTOInferVPTOVecScope.cpp b/lib/PTO/Transforms/PTOInferVPTOVecScope.cpp index 28d291a610..834fcd6ec8 100644 --- a/lib/PTO/Transforms/PTOInferVPTOVecScope.cpp +++ b/lib/PTO/Transforms/PTOInferVPTOVecScope.cpp @@ -89,7 +89,11 @@ static bool isExplicitVectorScopeCarrier(Operation *op) { } static bool isForbiddenInsideInferredVectorScope(Operation *op) { - return isa(op); + // Bisheng cannot expand block-query results produced inside an AIV vector + // scope. Keep these scalar queries outside the inferred scope and capture + // their results instead. + return isa(op); } static bool isVectorScopeBoundaryOperation(Operation *op) { diff --git a/test/lit/vpto/auto_vecscope_infer_block_query_boundary.pto b/test/lit/vpto/auto_vecscope_infer_block_query_boundary.pto new file mode 100644 index 0000000000..acaac5f4c7 --- /dev/null +++ b/test/lit/vpto/auto_vecscope_infer_block_query_boundary.pto @@ -0,0 +1,72 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --pto-level=level3 --emit-vpto %s -o - | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "vpto", pto.kernel_kind = #pto.kernel_kind} { + func.func @auto_vecscope_infer_block_query_boundary() { + %c0_i64 = arith.constant 0 : i64 + %ub = pto.castptr %c0_i64 : i64 -> !pto.ptr + %c0 = arith.constant 0 : index + %c6 = arith.constant 6 : index + %c1 = arith.constant 1 : index + %c8_i64 = arith.constant 8 : i64 + %c43 = arith.constant 43 : index + + scf.for %iv = %c0 to %c6 step %c1 { + %block_idx = pto.get_block_idx + %quotient = arith.floordivsi %block_idx, %c8_i64 : i64 + %index = arith.index_cast %quotient : i64 to index + %in_bounds = arith.cmpi slt, %index, %c43 : index + scf.if %in_bounds { + %value = pto.vlds %ub[%c0] : !pto.ptr -> !pto.vreg<128xbf16> + %mask = pto.pset_b16 "PAT_ALL" : !pto.mask + pto.vsts %value, %ub[%c0], %mask : !pto.vreg<128xbf16>, !pto.ptr, !pto.mask + } + } + return + } + + func.func @auto_vecscope_infer_block_num_boundary() { + %c0_i64 = arith.constant 0 : i64 + %ub = pto.castptr %c0_i64 : i64 -> !pto.ptr + %c0 = arith.constant 0 : index + %c1_i64 = arith.constant 1 : i64 + %block_num = pto.get_block_num + %has_blocks = arith.cmpi sge, %block_num, %c1_i64 : i64 + scf.if %has_blocks { + %value = pto.vlds %ub[%c0] : !pto.ptr -> !pto.vreg<128xbf16> + %mask = pto.pset_b16 "PAT_ALL" : !pto.mask + pto.vsts %value, %ub[%c0], %mask : !pto.vreg<128xbf16>, !pto.ptr, !pto.mask + } + return + } + } +} + +// CHECK-LABEL: func.func @auto_vecscope_infer_block_query_boundary +// CHECK: %[[BLOCK_IDX:.*]] = pto.get_block_idx +// CHECK: scf.for +// CHECK-NEXT: pto.vecscope { +// CHECK-NOT: pto.get_block_idx +// CHECK: arith.floordivsi %[[BLOCK_IDX]] +// CHECK: scf.if +// CHECK: pto.vlds +// CHECK: pto.vsts +// CHECK: } + +// CHECK-LABEL: func.func @auto_vecscope_infer_block_num_boundary +// CHECK: %[[BLOCK_NUM:.*]] = pto.get_block_num +// CHECK-NEXT: pto.vecscope { +// CHECK-NOT: pto.get_block_num +// CHECK: arith.cmpi sge, %[[BLOCK_NUM]] +// CHECK: scf.if +// CHECK: pto.vlds +// CHECK: pto.vsts +// CHECK: } From de8e4cef727c40be1f14631f5559b0c51e7efe70 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 09:32:01 +0800 Subject: [PATCH 029/122] Add nightly wheel installation support --- .github/workflows/build_wheel.yml | 2 +- .github/workflows/build_wheel_mac.yml | 2 +- CMakeLists.txt | 9 ++ README.md | 15 +++ test/python/install_nightly_wheel.py | 60 ++++++++++ tools/install_nightly_wheel.py | 161 ++++++++++++++++++++++++++ 6 files changed, 247 insertions(+), 2 deletions(-) create mode 100644 test/python/install_nightly_wheel.py create mode 100644 tools/install_nightly_wheel.py diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml index 1efe47f612..1b19385df7 100644 --- a/.github/workflows/build_wheel.yml +++ b/.github/workflows/build_wheel.yml @@ -32,7 +32,7 @@ permissions: contents: read concurrency: - group: build-wheel-${{ github.event.pull_request.number || github.ref }} + group: ${{ github.event_name == 'schedule' && 'nightly-wheel-publish' || format('build-wheel-{0}', github.event.pull_request.number || github.ref) }} cancel-in-progress: ${{ github.event_name == 'push' || github.event_name == 'pull_request' }} env: diff --git a/.github/workflows/build_wheel_mac.yml b/.github/workflows/build_wheel_mac.yml index f9afd60030..d7e217a522 100644 --- a/.github/workflows/build_wheel_mac.yml +++ b/.github/workflows/build_wheel_mac.yml @@ -33,7 +33,7 @@ permissions: contents: read concurrency: - group: build-wheel-macos-${{ github.event.pull_request.number || github.ref }} + group: ${{ github.event_name == 'schedule' && 'nightly-wheel-publish' || format('build-wheel-macos-{0}', github.event.pull_request.number || github.ref) }} cancel-in-progress: ${{ github.event_name == 'push' || github.event_name == 'pull_request' }} env: diff --git a/CMakeLists.txt b/CMakeLists.txt index ef73e3e15c..119125477d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -237,6 +237,15 @@ if(BUILD_TESTING) LABELS "PTO" ENVIRONMENT "${_pto_python_test_env}" ) + add_test( + NAME pto_install_nightly_wheel + COMMAND "${Python3_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/test/python/install_nightly_wheel.py" + ) + set_tests_properties(pto_install_nightly_wheel PROPERTIES + LABELS "PTO" + ENVIRONMENT "${_pto_python_test_env}" + ) endif() add_subdirectory(tools/ptobc/tests) endif() diff --git a/README.md b/README.md index 337d4f099a..abe1cfc292 100644 --- a/README.md +++ b/README.md @@ -215,6 +215,21 @@ ninja -C "$PTO_SOURCE_DIR/build" check-pto 发布 wheel 自带运行时依赖，不使用外部 LLVM build tree 时无需设置上述 `LD_LIBRARY_PATH`。无论哪种安装方式，都不需要手工拼接 `PYTHONPATH`。 +### Daily wheel + +定时构建会将最新 wheel 发布到 GitHub 的 `nightly` release。开发者可以查看 +[Nightly Build](https://github.com/hw-native-sys/PTOAS/releases/tag/nightly)， +或在仓库 checkout 中运行下面的命令自动选择当前 Python 和平台对应的 wheel： + +```bash +python tools/install_nightly_wheel.py +``` + +目前 daily workflow 提供 Python 3.10、3.11、3.12 的 Linux 和 macOS wheel。 + +如需先查看将要安装的文件，可以加上 `--dry-run`。脚本使用当前 Python +环境执行安装，并会替换该环境中已安装的同名 nightly wheel。 + 需要 CANN、Bisheng、simulator 或 NPU 时，再加载 CANN 对外提供的环境脚本。 常见安装位置如下，按实际环境选择一个： diff --git a/test/python/install_nightly_wheel.py b/test/python/install_nightly_wheel.py new file mode 100644 index 0000000000..6da9b9a033 --- /dev/null +++ b/test/python/install_nightly_wheel.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See the License for details. + +import importlib.util +import unittest +from pathlib import Path +from unittest import mock + +from packaging.tags import Tag + + +SCRIPT = Path(__file__).resolve().parents[2] / "tools" / "install_nightly_wheel.py" +SPEC = importlib.util.spec_from_file_location("install_nightly_wheel", SCRIPT) +assert SPEC and SPEC.loader +INSTALLER = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(INSTALLER) + + +class NightlyWheelSelectionTests(unittest.TestCase): + def test_selects_latest_compatible_version(self): + compatible = Tag("cp312", "cp312", "manylinux_2_34_x86_64") + release = { + "tag_name": "nightly", + "assets": [ + { + "name": "ptoas-0.56-cp312-cp312-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/old.whl", + }, + { + "name": "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/new.whl", + }, + { + "name": "ptoas-0.58-cp311-cp311-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/wrong-python.whl", + }, + ], + } + + with mock.patch("packaging.tags.sys_tags", return_value=iter([compatible])): + name, url = INSTALLER.select_wheel(release, "ptoas") + + self.assertEqual(name, "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl") + self.assertEqual(url, "https://example.invalid/new.whl") + + def test_rejects_missing_compatible_wheel(self): + release = {"tag_name": "nightly", "assets": []} + with self.assertRaisesRegex(RuntimeError, "no compatible ptoas wheel"): + with mock.patch("packaging.tags.sys_tags", return_value=iter(())): + INSTALLER.select_wheel(release, "ptoas") + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/install_nightly_wheel.py b/tools/install_nightly_wheel.py new file mode 100644 index 0000000000..37328cbf32 --- /dev/null +++ b/tools/install_nightly_wheel.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See the License for details. + +"""Install the latest wheel published by the PTOAS nightly GitHub release.""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import tempfile +import urllib.error +import urllib.request +from pathlib import Path + + +DEFAULT_REPOSITORY = "hw-native-sys/PTOAS" +DEFAULT_TAG = "nightly" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Install the latest compatible wheel from a PTOAS GitHub release." + ) + parser.add_argument( + "--repository", + default=DEFAULT_REPOSITORY, + help=f"GitHub repository (default: {DEFAULT_REPOSITORY})", + ) + parser.add_argument( + "--tag", + default=DEFAULT_TAG, + help=f"GitHub release tag (default: {DEFAULT_TAG})", + ) + parser.add_argument( + "--package", + default="ptoas", + help="Distribution to install (default: ptoas)", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Resolve and print the wheel without installing it", + ) + return parser.parse_args() + + +def github_request(url: str) -> object: + headers = { + "Accept": "application/vnd.github+json", + "User-Agent": "ptoas-nightly-wheel-installer", + } + token = os.environ.get("GITHUB_TOKEN") + if token: + headers["Authorization"] = f"Bearer {token}" + request = urllib.request.Request(url, headers=headers) + try: + with urllib.request.urlopen(request) as response: + return json.load(response) + except urllib.error.HTTPError as error: + raise RuntimeError( + f"GitHub API request failed with HTTP {error.code}: {error.reason}" + ) from error + except urllib.error.URLError as error: + raise RuntimeError(f"unable to reach GitHub API: {error.reason}") from error + + +def select_wheel(release: object, package: str) -> tuple[str, str]: + try: + from packaging.tags import sys_tags + from packaging.utils import parse_wheel_filename + except ImportError as error: + raise RuntimeError( + "the packaging module is required; install it with 'python -m pip install packaging'" + ) from error + + if not isinstance(release, dict) or not isinstance(release.get("assets"), list): + raise RuntimeError("GitHub release response does not contain wheel assets") + + supported_tags = list(sys_tags()) + tag_rank = {tag: rank for rank, tag in enumerate(supported_tags)} + candidates = [] + for asset in release["assets"]: + if not isinstance(asset, dict): + continue + name = asset.get("name") + url = asset.get("browser_download_url") + if not isinstance(name, str) or not name.endswith(".whl") or not isinstance(url, str): + continue + try: + distribution, version, _, wheel_tags = parse_wheel_filename(name) + except (TypeError, ValueError): + continue + if str(distribution) != package.replace("_", "-").lower(): + continue + matching_ranks = [tag_rank[tag] for tag in wheel_tags if tag in tag_rank] + if matching_ranks: + candidates.append((version, min(matching_ranks), name, url)) + + if not candidates: + raise RuntimeError( + f"no compatible {package} wheel found in the {release.get('tag_name', 'requested')} release" + ) + _, _, name, url = max(candidates, key=lambda item: (item[0], -item[1], item[2])) + return name, url + + +def download(url: str, destination: Path) -> None: + request = urllib.request.Request( + url, + headers={"User-Agent": "ptoas-nightly-wheel-installer"}, + ) + try: + with urllib.request.urlopen(request) as response, destination.open("wb") as output: + while chunk := response.read(1024 * 1024): + output.write(chunk) + except (OSError, urllib.error.URLError) as error: + raise RuntimeError(f"failed to download wheel: {error}") from error + + +def main() -> int: + args = parse_args() + try: + release_url = ( + f"https://api.github.com/repos/{args.repository}/releases/tags/{args.tag}" + ) + release = github_request(release_url) + wheel_name, wheel_url = select_wheel(release, args.package) + print(f"Selected wheel: {wheel_name}") + if args.dry_run: + print(wheel_url) + return 0 + + with tempfile.TemporaryDirectory(prefix="ptoas-nightly-") as directory: + wheel_path = Path(directory) / wheel_name + download(wheel_url, wheel_path) + command = [ + sys.executable, + "-m", + "pip", + "install", + "--force-reinstall", + str(wheel_path), + ] + subprocess.run(command, check=True) + except (RuntimeError, subprocess.CalledProcessError) as error: + print(f"error: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 46fa5062672927619cf5b953762a7cc276695e16 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 09:41:55 +0800 Subject: [PATCH 030/122] Fix nightly wheel license headers --- test/python/install_nightly_wheel.py | 4 ++-- tools/install_nightly_wheel.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/test/python/install_nightly_wheel.py b/test/python/install_nightly_wheel.py index 6da9b9a033..29d5103ae3 100644 --- a/test/python/install_nightly_wheel.py +++ b/test/python/install_nightly_wheel.py @@ -2,10 +2,10 @@ # Copyright (c) 2026 Huawei Technologies Co., Ltd. # This program is free software, you can redistribute it and/or modify it under the terms and conditions of # CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for details. +# Please refer to the License for details. You may not use this file except in compliance with the License. # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -# See the License for details. +# See LICENSE in the root of the software repository for the full text of the License. import importlib.util import unittest diff --git a/tools/install_nightly_wheel.py b/tools/install_nightly_wheel.py index 37328cbf32..0828b632f8 100644 --- a/tools/install_nightly_wheel.py +++ b/tools/install_nightly_wheel.py @@ -2,10 +2,10 @@ # Copyright (c) 2026 Huawei Technologies Co., Ltd. # This program is free software, you can redistribute it and/or modify it under the terms and conditions of # CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for details. +# Please refer to the License for details. You may not use this file except in compliance with the License. # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -# See the License for details. +# See LICENSE in the root of the software repository for the full text of the License. """Install the latest wheel published by the PTOAS nightly GitHub release.""" From 168ab692da8b50cd24ed03b338f6133d0489fb87 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 10:33:54 +0800 Subject: [PATCH 031/122] Address nightly wheel review comments --- .github/workflows/build_wheel.yml | 5 + .github/workflows/build_wheel_mac.yml | 5 + CMakeLists.txt | 2 +- README.md | 8 +- test/python/install_nightly_wheel.py | 60 -------- test/python/test_install_nightly_wheel.py | 177 ++++++++++++++++++++++ tools/install_nightly_wheel.py | 126 +++++++++++++-- 7 files changed, 304 insertions(+), 79 deletions(-) delete mode 100644 test/python/install_nightly_wheel.py create mode 100644 test/python/test_install_nightly_wheel.py diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml index 1b19385df7..22c01e1557 100644 --- a/.github/workflows/build_wheel.yml +++ b/.github/workflows/build_wheel.yml @@ -32,6 +32,11 @@ permissions: contents: read concurrency: + # Linux and macOS publish to the same rolling release. The schedules are + # offset by 30 minutes; this non-canceling group prevents asset upload races. + # GitHub keeps at most one pending run in a group, so a run that spans the + # next daily schedule may replace an older pending run by design. Manual + # workflow_dispatch runs are dry-runs and do not publish the nightly release. group: ${{ github.event_name == 'schedule' && 'nightly-wheel-publish' || format('build-wheel-{0}', github.event.pull_request.number || github.ref) }} cancel-in-progress: ${{ github.event_name == 'push' || github.event_name == 'pull_request' }} diff --git a/.github/workflows/build_wheel_mac.yml b/.github/workflows/build_wheel_mac.yml index d7e217a522..f4698a8503 100644 --- a/.github/workflows/build_wheel_mac.yml +++ b/.github/workflows/build_wheel_mac.yml @@ -33,6 +33,11 @@ permissions: contents: read concurrency: + # Linux and macOS publish to the same rolling release. The schedules are + # offset by 30 minutes; this non-canceling group prevents asset upload races. + # GitHub keeps at most one pending run in a group, so a run that spans the + # next daily schedule may replace an older pending run by design. Manual + # workflow_dispatch runs are dry-runs and do not publish the nightly release. group: ${{ github.event_name == 'schedule' && 'nightly-wheel-publish' || format('build-wheel-macos-{0}', github.event.pull_request.number || github.ref) }} cancel-in-progress: ${{ github.event_name == 'push' || github.event_name == 'pull_request' }} diff --git a/CMakeLists.txt b/CMakeLists.txt index 119125477d..13ea3accf0 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -240,7 +240,7 @@ if(BUILD_TESTING) add_test( NAME pto_install_nightly_wheel COMMAND "${Python3_EXECUTABLE}" - "${CMAKE_CURRENT_SOURCE_DIR}/test/python/install_nightly_wheel.py" + "${CMAKE_CURRENT_SOURCE_DIR}/test/python/test_install_nightly_wheel.py" ) set_tests_properties(pto_install_nightly_wheel PROPERTIES LABELS "PTO" diff --git a/README.md b/README.md index abe1cfc292..243fda727f 100644 --- a/README.md +++ b/README.md @@ -225,10 +225,14 @@ ninja -C "$PTO_SOURCE_DIR/build" check-pto python tools/install_nightly_wheel.py ``` -目前 daily workflow 提供 Python 3.10、3.11、3.12 的 Linux 和 macOS wheel。 +daily workflow 的实际 Python 版本、平台和架构以 nightly release 中当前发布的 +wheel 为准。 如需先查看将要安装的文件，可以加上 `--dry-run`。脚本使用当前 Python -环境执行安装，并会替换该环境中已安装的同名 nightly wheel。 +环境执行安装，不会自动重装 wheel 的运行时依赖，并会替换该环境中已安装的同名 +nightly wheel。GitHub Release 提供 asset digest 时脚本会自动校验 SHA-256，也可 +通过 `--sha256` 显式指定摘要。nightly wheel 来自 GitHub Release，使用前请确认 +下载来源和当前环境符合预期。若选中的 asset 超过 48 小时未更新，脚本会给出警告。 需要 CANN、Bisheng、simulator 或 NPU 时，再加载 CANN 对外提供的环境脚本。 常见安装位置如下，按实际环境选择一个： diff --git a/test/python/install_nightly_wheel.py b/test/python/install_nightly_wheel.py deleted file mode 100644 index 29d5103ae3..0000000000 --- a/test/python/install_nightly_wheel.py +++ /dev/null @@ -1,60 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software, you can redistribute it and/or modify it under the terms and conditions of -# CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -# See LICENSE in the root of the software repository for the full text of the License. - -import importlib.util -import unittest -from pathlib import Path -from unittest import mock - -from packaging.tags import Tag - - -SCRIPT = Path(__file__).resolve().parents[2] / "tools" / "install_nightly_wheel.py" -SPEC = importlib.util.spec_from_file_location("install_nightly_wheel", SCRIPT) -assert SPEC and SPEC.loader -INSTALLER = importlib.util.module_from_spec(SPEC) -SPEC.loader.exec_module(INSTALLER) - - -class NightlyWheelSelectionTests(unittest.TestCase): - def test_selects_latest_compatible_version(self): - compatible = Tag("cp312", "cp312", "manylinux_2_34_x86_64") - release = { - "tag_name": "nightly", - "assets": [ - { - "name": "ptoas-0.56-cp312-cp312-manylinux_2_34_x86_64.whl", - "browser_download_url": "https://example.invalid/old.whl", - }, - { - "name": "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl", - "browser_download_url": "https://example.invalid/new.whl", - }, - { - "name": "ptoas-0.58-cp311-cp311-manylinux_2_34_x86_64.whl", - "browser_download_url": "https://example.invalid/wrong-python.whl", - }, - ], - } - - with mock.patch("packaging.tags.sys_tags", return_value=iter([compatible])): - name, url = INSTALLER.select_wheel(release, "ptoas") - - self.assertEqual(name, "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl") - self.assertEqual(url, "https://example.invalid/new.whl") - - def test_rejects_missing_compatible_wheel(self): - release = {"tag_name": "nightly", "assets": []} - with self.assertRaisesRegex(RuntimeError, "no compatible ptoas wheel"): - with mock.patch("packaging.tags.sys_tags", return_value=iter(())): - INSTALLER.select_wheel(release, "ptoas") - - -if __name__ == "__main__": - unittest.main() diff --git a/test/python/test_install_nightly_wheel.py b/test/python/test_install_nightly_wheel.py new file mode 100644 index 0000000000..30018ab42c --- /dev/null +++ b/test/python/test_install_nightly_wheel.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import importlib.util +import hashlib +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +try: + from packaging.tags import Tag +except ImportError: # pragma: no cover - exercised only in minimal build environments. + Tag = None + + +SCRIPT = Path(__file__).resolve().parents[2] / "tools" / "install_nightly_wheel.py" +SPEC = importlib.util.spec_from_file_location("install_nightly_wheel", SCRIPT) +assert SPEC and SPEC.loader +INSTALLER = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(INSTALLER) + + +@unittest.skipIf(Tag is None, "packaging is not installed") +class NightlyWheelSelectionTests(unittest.TestCase): + def test_selects_latest_compatible_version(self): + compatible = Tag("cp312", "cp312", "manylinux_2_34_x86_64") + release = { + "tag_name": "nightly", + "assets": [ + { + "name": "ptoas-0.56-cp312-cp312-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/old.whl", + "updated_at": "2026-08-06T00:00:00Z", + }, + { + "name": "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/new.whl", + "updated_at": "2026-08-06T01:00:00Z", + }, + { + "name": "ptoas-0.58-cp311-cp311-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/wrong-python.whl", + }, + ], + } + + with mock.patch("packaging.tags.sys_tags", return_value=iter([compatible])): + selection = INSTALLER.select_wheel(release) + + self.assertEqual(selection.name, "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl") + self.assertEqual(selection.url, "https://example.invalid/new.whl") + + def test_prefers_newer_asset_over_higher_version(self): + compatible = Tag("cp312", "cp312", "manylinux_2_34_x86_64") + release = { + "tag_name": "nightly", + "assets": [ + { + "name": "ptoas-0.58-cp312-cp312-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/old.whl", + "updated_at": "2026-08-05T00:00:00Z", + }, + { + "name": "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/new.whl", + "updated_at": "2026-08-06T01:00:00Z", + }, + ], + } + + with mock.patch("packaging.tags.sys_tags", return_value=iter([compatible])): + selection = INSTALLER.select_wheel(release, "ptoas") + + self.assertEqual(selection.name, "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl") + + def test_prefers_higher_build_number(self): + compatible = Tag("cp312", "cp312", "manylinux_2_34_x86_64") + release = { + "tag_name": "nightly", + "assets": [ + { + "name": "ptoas-0.57-1-cp312-cp312-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/build1.whl", + "updated_at": "2026-08-06T01:00:00Z", + }, + { + "name": "ptoas-0.57-2-cp312-cp312-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/build2.whl", + "updated_at": "2026-08-06T01:00:00Z", + }, + ], + } + + with mock.patch("packaging.tags.sys_tags", return_value=iter([compatible])): + selection = INSTALLER.select_wheel(release, "ptoas") + + self.assertEqual(selection.name, "ptoas-0.57-2-cp312-cp312-manylinux_2_34_x86_64.whl") + + def test_normalizes_distribution_name(self): + compatible = Tag("cp312", "cp312", "manylinux_2_34_x86_64") + release = { + "tag_name": "nightly", + "assets": [ + { + "name": "pto_as-0.57-cp312-cp312-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/pto-as.whl", + } + ], + } + + with mock.patch("packaging.tags.sys_tags", return_value=iter([compatible])): + selection = INSTALLER.select_wheel(release, "pto_as") + + self.assertEqual(selection.name, "pto_as-0.57-cp312-cp312-manylinux_2_34_x86_64.whl") + + def test_prefers_better_supported_tag(self): + platform_tag = Tag("cp312", "cp312", "manylinux_2_34_x86_64") + universal_tag = Tag("py3", "none", "any") + release = { + "tag_name": "nightly", + "assets": [ + { + "name": "ptoas-0.57-py3-none-any.whl", + "browser_download_url": "https://example.invalid/universal.whl", + "updated_at": "2026-08-06T01:00:00Z", + }, + { + "name": "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/platform.whl", + "updated_at": "2026-08-06T01:00:00Z", + }, + ], + } + + with mock.patch("packaging.tags.sys_tags", return_value=iter([platform_tag, universal_tag])): + selection = INSTALLER.select_wheel(release, "ptoas") + + self.assertEqual(selection.name, "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl") + + def test_rejects_missing_compatible_wheel(self): + release = {"tag_name": "nightly", "assets": []} + with self.assertRaisesRegex(RuntimeError, "no compatible ptoas wheel"): + with mock.patch("packaging.tags.sys_tags", return_value=iter(())): + INSTALLER.select_wheel(release, "ptoas") + + def test_download_verifies_sha256(self): + payload = b"nightly wheel" + expected = hashlib.sha256(payload).hexdigest() + response = mock.MagicMock() + response.__enter__.return_value = response + response.read.side_effect = [payload, b""] + with tempfile.TemporaryDirectory() as directory: + destination = Path(directory) / "wheel.whl" + with mock.patch("urllib.request.urlopen", return_value=response): + INSTALLER.download("https://example.invalid/wheel.whl", destination, expected) + self.assertEqual(destination.read_bytes(), payload) + + def test_download_rejects_wrong_sha256(self): + response = mock.MagicMock() + response.__enter__.return_value = response + response.read.side_effect = [b"nightly wheel", b""] + with tempfile.TemporaryDirectory() as directory: + destination = Path(directory) / "wheel.whl" + with mock.patch("urllib.request.urlopen", return_value=response): + with self.assertRaisesRegex(RuntimeError, "SHA-256 mismatch"): + INSTALLER.download("https://example.invalid/wheel.whl", destination, "0" * 64) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/install_nightly_wheel.py b/tools/install_nightly_wheel.py index 0828b632f8..39236c7518 100644 --- a/tools/install_nightly_wheel.py +++ b/tools/install_nightly_wheel.py @@ -12,6 +12,8 @@ from __future__ import annotations import argparse +import datetime +import hashlib import json import os import subprocess @@ -24,6 +26,22 @@ DEFAULT_REPOSITORY = "hw-native-sys/PTOAS" DEFAULT_TAG = "nightly" +NETWORK_TIMEOUT_SECONDS = 30 +STALE_WHEEL_AGE = datetime.timedelta(hours=48) + + +class WheelSelection: + def __init__( + self, + name: str, + url: str, + updated_at: datetime.datetime | None, + digest: str | None, + ) -> None: + self.name = name + self.url = url + self.updated_at = updated_at + self.digest = digest def parse_args() -> argparse.Namespace: @@ -50,6 +68,10 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Resolve and print the wheel without installing it", ) + parser.add_argument( + "--sha256", + help="Expected SHA-256 digest of the selected wheel", + ) return parser.parse_args() @@ -63,20 +85,53 @@ def github_request(url: str) -> object: headers["Authorization"] = f"Bearer {token}" request = urllib.request.Request(url, headers=headers) try: - with urllib.request.urlopen(request) as response: + with urllib.request.urlopen(request, timeout=NETWORK_TIMEOUT_SECONDS) as response: return json.load(response) except urllib.error.HTTPError as error: + detail = "" + if error.code in (403, 429): + detail = " Set GITHUB_TOKEN if GitHub API rate limiting is suspected." + try: + response_body = error.read().decode("utf-8", errors="replace") + response_json = json.loads(response_body) + message = response_json.get("message") if isinstance(response_json, dict) else None + if message: + detail += f" GitHub message: {message}." + except (OSError, UnicodeError, json.JSONDecodeError): + pass + raise RuntimeError( + f"GitHub API request failed with HTTP {error.code}: {error.reason}.{detail}" + ) from error + except TimeoutError as error: raise RuntimeError( - f"GitHub API request failed with HTTP {error.code}: {error.reason}" + f"GitHub API request timed out after {NETWORK_TIMEOUT_SECONDS} seconds" ) from error except urllib.error.URLError as error: raise RuntimeError(f"unable to reach GitHub API: {error.reason}") from error -def select_wheel(release: object, package: str) -> tuple[str, str]: +def parse_updated_at(value: object) -> datetime.datetime | None: + if not isinstance(value, str): + return None + try: + return datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + + +def parse_digest(value: object) -> str | None: + if not isinstance(value, str): + return None + digest = value.removeprefix("sha256:") + if len(digest) == 64 and all(character in "0123456789abcdefABCDEF" for character in digest): + return digest.lower() + return None + + +def select_wheel(release: object, package: str = "ptoas") -> WheelSelection: try: from packaging.tags import sys_tags - from packaging.utils import parse_wheel_filename + from packaging.utils import canonicalize_name, parse_wheel_filename except ImportError as error: raise RuntimeError( "the packaging module is required; install it with 'python -m pip install packaging'" @@ -87,6 +142,7 @@ def select_wheel(release: object, package: str) -> tuple[str, str]: supported_tags = list(sys_tags()) tag_rank = {tag: rank for rank, tag in enumerate(supported_tags)} + canonical_package = canonicalize_name(package) candidates = [] for asset in release["assets"]: if not isinstance(asset, dict): @@ -96,34 +152,63 @@ def select_wheel(release: object, package: str) -> tuple[str, str]: if not isinstance(name, str) or not name.endswith(".whl") or not isinstance(url, str): continue try: - distribution, version, _, wheel_tags = parse_wheel_filename(name) + distribution, version, build, wheel_tags = parse_wheel_filename(name) except (TypeError, ValueError): continue - if str(distribution) != package.replace("_", "-").lower(): + if canonicalize_name(str(distribution)) != canonical_package: continue matching_ranks = [tag_rank[tag] for tag in wheel_tags if tag in tag_rank] if matching_ranks: - candidates.append((version, min(matching_ranks), name, url)) + candidates.append( + ( + parse_updated_at(asset.get("updated_at")), + version, + build or (-1, ""), + min(matching_ranks), + name, + url, + parse_digest(asset.get("digest")), + ) + ) if not candidates: raise RuntimeError( f"no compatible {package} wheel found in the {release.get('tag_name', 'requested')} release" ) - _, _, name, url = max(candidates, key=lambda item: (item[0], -item[1], item[2])) - return name, url + selected = max( + candidates, + key=lambda item: ( + item[0] or datetime.datetime.min.replace(tzinfo=datetime.timezone.utc), + item[1], + item[2], + -item[3], + item[4], + ), + ) + return WheelSelection(selected[4], selected[5], selected[0], selected[6]) -def download(url: str, destination: Path) -> None: +def download(url: str, destination: Path, expected_sha256: str | None = None) -> None: request = urllib.request.Request( url, headers={"User-Agent": "ptoas-nightly-wheel-installer"}, ) try: - with urllib.request.urlopen(request) as response, destination.open("wb") as output: + with urllib.request.urlopen(request, timeout=NETWORK_TIMEOUT_SECONDS) as response, destination.open( + "wb" + ) as output: + digest = hashlib.sha256() while chunk := response.read(1024 * 1024): output.write(chunk) + digest.update(chunk) + except TimeoutError as error: + raise RuntimeError( + f"wheel download timed out after {NETWORK_TIMEOUT_SECONDS} seconds" + ) from error except (OSError, urllib.error.URLError) as error: raise RuntimeError(f"failed to download wheel: {error}") from error + if expected_sha256 and digest.hexdigest() != expected_sha256.lower().removeprefix("sha256:"): + raise RuntimeError(f"SHA-256 mismatch for downloaded wheel {destination.name}") def main() -> int: @@ -133,21 +218,30 @@ def main() -> int: f"https://api.github.com/repos/{args.repository}/releases/tags/{args.tag}" ) release = github_request(release_url) - wheel_name, wheel_url = select_wheel(release, args.package) - print(f"Selected wheel: {wheel_name}") + selection = select_wheel(release, args.package) + print(f"Selected wheel: {selection.name}") + if selection.updated_at: + age = datetime.datetime.now(datetime.timezone.utc) - selection.updated_at + if age > STALE_WHEEL_AGE: + print( + f"warning: selected wheel was last updated {age.total_seconds() / 3600:.1f} hours ago", + file=sys.stderr, + ) if args.dry_run: - print(wheel_url) + print(selection.url) return 0 with tempfile.TemporaryDirectory(prefix="ptoas-nightly-") as directory: - wheel_path = Path(directory) / wheel_name - download(wheel_url, wheel_path) + wheel_path = Path(directory) / selection.name + expected_sha256 = args.sha256 or selection.digest + download(selection.url, wheel_path, expected_sha256) command = [ sys.executable, "-m", "pip", "install", "--force-reinstall", + "--no-deps", str(wheel_path), ] subprocess.run(command, check=True) From df44ec0fab57ed1a7a6f4fdac88d740831c0a593 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 10:35:01 +0800 Subject: [PATCH 032/122] Handle timezone-less release timestamps --- tools/install_nightly_wheel.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tools/install_nightly_wheel.py b/tools/install_nightly_wheel.py index 39236c7518..19329a93f6 100644 --- a/tools/install_nightly_wheel.py +++ b/tools/install_nightly_wheel.py @@ -114,7 +114,10 @@ def parse_updated_at(value: object) -> datetime.datetime | None: if not isinstance(value, str): return None try: - return datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + timestamp = datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + if timestamp.tzinfo is None: + timestamp = timestamp.replace(tzinfo=datetime.timezone.utc) + return timestamp except ValueError: return None From 619c1cd2c2a48ce685ae44c0e487cdf835202a5f Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 19:34:25 +0800 Subject: [PATCH 033/122] Make nightly installer self-contained --- .github/workflows/build_wheel.yml | 41 +++++++++++++++++++++ .github/workflows/build_wheel_mac.yml | 41 +++++++++++++++++++++ README.md | 5 +++ test/python/test_install_nightly_wheel.py | 43 +++++++++++++++++++++-- tools/install_nightly_wheel.py | 22 ++++++++++-- 5 files changed, 146 insertions(+), 6 deletions(-) diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml index 22c01e1557..8878a7c8fb 100644 --- a/.github/workflows/build_wheel.yml +++ b/.github/workflows/build_wheel.yml @@ -409,6 +409,34 @@ jobs: mv "release-artifacts/ptoas-bin-aarch64/ptoas-bin-aarch64.tar.gz" \ "release-artifacts/ptoas-bin-aarch64.tar.gz" + - name: Create nightly build manifest + if: github.event_name == 'schedule' + env: + MANIFEST_PATH: release-artifacts/nightly-manifest-linux.json + run: | + set -euo pipefail + python3 - <<'PY' + import json + import os + from datetime import datetime, timezone + from pathlib import Path + + root = Path("release-artifacts") + assets = sorted(path.name for path in root.iterdir() if path.is_file()) + manifest = { + "platform": "linux", + "source_commit": os.environ["GITHUB_SHA"], + "workflow": os.environ["GITHUB_WORKFLOW"], + "run_id": os.environ["GITHUB_RUN_ID"], + "run_url": f"{os.environ['GITHUB_SERVER_URL']}/{os.environ['GITHUB_REPOSITORY']}/actions/runs/{os.environ['GITHUB_RUN_ID']}", + "generated_at_utc": datetime.now(timezone.utc).isoformat(), + "assets": assets, + } + Path(os.environ["MANIFEST_PATH"]).write_text( + json.dumps(manifest, indent=2) + "\n", encoding="utf-8" + ) + PY + - name: Upload wheel assets to GitHub Release uses: softprops/action-gh-release@v2 with: @@ -433,6 +461,19 @@ jobs: overwrite_files: true files: release-artifacts/*.tar.gz + - name: Upload nightly build manifest + if: github.event_name == 'schedule' + uses: softprops/action-gh-release@v2 + with: + tag_name: ${{ env.RELEASE_TAG }} + name: ${{ env.RELEASE_NAME }} + target_commitish: ${{ github.sha }} + prerelease: ${{ env.RELEASE_PRERELEASE }} + make_latest: ${{ env.RELEASE_MAKE_LATEST }} + body: ${{ env.RELEASE_BODY }} + overwrite_files: true + files: release-artifacts/nightly-manifest-linux.json + bump_base_version: name: Bump base version after release if: github.event_name == 'release' && github.event.action == 'released' && (startsWith(github.ref_name, 'ptoas-v') || (startsWith(github.ref_name, 'v') && !startsWith(github.ref_name, 'vmi-v'))) diff --git a/.github/workflows/build_wheel_mac.yml b/.github/workflows/build_wheel_mac.yml index f4698a8503..7b76eb13a7 100644 --- a/.github/workflows/build_wheel_mac.yml +++ b/.github/workflows/build_wheel_mac.yml @@ -485,6 +485,34 @@ jobs: mv "release-artifacts/ptoas-bin-macos-aarch64/ptoas-bin-macos-aarch64.tar.gz" \ "release-artifacts/ptoas-bin-macos-aarch64.tar.gz" + - name: Create nightly build manifest + if: github.event_name == 'schedule' + env: + MANIFEST_PATH: release-artifacts/nightly-manifest-macos.json + run: | + set -euo pipefail + python3 - <<'PY' + import json + import os + from datetime import datetime, timezone + from pathlib import Path + + root = Path("release-artifacts") + assets = sorted(path.name for path in root.iterdir() if path.is_file()) + manifest = { + "platform": "macos", + "source_commit": os.environ["GITHUB_SHA"], + "workflow": os.environ["GITHUB_WORKFLOW"], + "run_id": os.environ["GITHUB_RUN_ID"], + "run_url": f"{os.environ['GITHUB_SERVER_URL']}/{os.environ['GITHUB_REPOSITORY']}/actions/runs/{os.environ['GITHUB_RUN_ID']}", + "generated_at_utc": datetime.now(timezone.utc).isoformat(), + "assets": assets, + } + Path(os.environ["MANIFEST_PATH"]).write_text( + json.dumps(manifest, indent=2) + "\n", encoding="utf-8" + ) + PY + - name: Upload assets to GitHub Release uses: softprops/action-gh-release@v2 with: @@ -498,3 +526,16 @@ jobs: files: | release-artifacts/*.whl release-artifacts/*.tar.gz + + - name: Upload nightly build manifest + if: github.event_name == 'schedule' + uses: softprops/action-gh-release@v2 + with: + tag_name: ${{ env.RELEASE_TAG }} + name: ${{ env.RELEASE_NAME }} + target_commitish: ${{ github.sha }} + prerelease: ${{ env.RELEASE_PRERELEASE }} + make_latest: ${{ env.RELEASE_MAKE_LATEST }} + body: ${{ env.RELEASE_BODY }} + overwrite_files: true + files: release-artifacts/nightly-manifest-macos.json diff --git a/README.md b/README.md index 243fda727f..ee4578d843 100644 --- a/README.md +++ b/README.md @@ -225,6 +225,9 @@ ninja -C "$PTO_SOURCE_DIR/build" check-pto python tools/install_nightly_wheel.py ``` +脚本使用当前 Python 自带的 pip packaging 支持选择 wheel，无需预先单独安装 +`packaging`。 + daily workflow 的实际 Python 版本、平台和架构以 nightly release 中当前发布的 wheel 为准。 @@ -233,6 +236,8 @@ wheel 为准。 nightly wheel。GitHub Release 提供 asset digest 时脚本会自动校验 SHA-256，也可 通过 `--sha256` 显式指定摘要。nightly wheel 来自 GitHub Release，使用前请确认 下载来源和当前环境符合预期。若选中的 asset 超过 48 小时未更新，脚本会给出警告。 +Linux 和 macOS nightly release 同时附带 manifest，其中记录各平台 wheel 对应的 +源码 commit 和构建任务，跨平台版本不一致时可据此核对。 需要 CANN、Bisheng、simulator 或 NPU 时，再加载 CANN 对外提供的环境脚本。 常见安装位置如下，按实际环境选择一个： diff --git a/test/python/test_install_nightly_wheel.py b/test/python/test_install_nightly_wheel.py index 30018ab42c..151af862bb 100644 --- a/test/python/test_install_nightly_wheel.py +++ b/test/python/test_install_nightly_wheel.py @@ -7,17 +7,22 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -import importlib.util import hashlib +import importlib.util import tempfile +import sys import unittest from pathlib import Path from unittest import mock + try: from packaging.tags import Tag except ImportError: # pragma: no cover - exercised only in minimal build environments. - Tag = None + try: + from pip._vendor.packaging.tags import Tag + except ImportError: + Tag = None SCRIPT = Path(__file__).resolve().parents[2] / "tools" / "install_nightly_wheel.py" @@ -27,8 +32,35 @@ SPEC.loader.exec_module(INSTALLER) -@unittest.skipIf(Tag is None, "packaging is not installed") class NightlyWheelSelectionTests(unittest.TestCase): + @unittest.skipIf(Tag is None, "packaging is not available in this Python environment") + def test_falls_back_to_pip_vendor_packaging(self): + from pip._vendor.packaging.tags import Tag as PipTag + + compatible = PipTag("cp312", "cp312", "manylinux_2_34_x86_64") + release = { + "tag_name": "nightly", + "assets": [ + { + "name": "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl", + "browser_download_url": "https://example.invalid/new.whl", + } + ], + } + from pip._vendor.packaging import tags as pip_tags + + with mock.patch.dict( + sys.modules, + {"packaging": None, "packaging.tags": None, "packaging.utils": None}, + ): + with mock.patch.object( + pip_tags, "sys_tags", return_value=iter([compatible]) + ): + selection = INSTALLER.select_wheel(release) + + self.assertEqual(selection.name, "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl") + + @unittest.skipIf(Tag is None, "packaging is not available in this Python environment") def test_selects_latest_compatible_version(self): compatible = Tag("cp312", "cp312", "manylinux_2_34_x86_64") release = { @@ -57,6 +89,7 @@ def test_selects_latest_compatible_version(self): self.assertEqual(selection.name, "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl") self.assertEqual(selection.url, "https://example.invalid/new.whl") + @unittest.skipIf(Tag is None, "packaging is not available in this Python environment") def test_prefers_newer_asset_over_higher_version(self): compatible = Tag("cp312", "cp312", "manylinux_2_34_x86_64") release = { @@ -80,6 +113,7 @@ def test_prefers_newer_asset_over_higher_version(self): self.assertEqual(selection.name, "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl") + @unittest.skipIf(Tag is None, "packaging is not available in this Python environment") def test_prefers_higher_build_number(self): compatible = Tag("cp312", "cp312", "manylinux_2_34_x86_64") release = { @@ -103,6 +137,7 @@ def test_prefers_higher_build_number(self): self.assertEqual(selection.name, "ptoas-0.57-2-cp312-cp312-manylinux_2_34_x86_64.whl") + @unittest.skipIf(Tag is None, "packaging is not available in this Python environment") def test_normalizes_distribution_name(self): compatible = Tag("cp312", "cp312", "manylinux_2_34_x86_64") release = { @@ -120,6 +155,7 @@ def test_normalizes_distribution_name(self): self.assertEqual(selection.name, "pto_as-0.57-cp312-cp312-manylinux_2_34_x86_64.whl") + @unittest.skipIf(Tag is None, "packaging is not available in this Python environment") def test_prefers_better_supported_tag(self): platform_tag = Tag("cp312", "cp312", "manylinux_2_34_x86_64") universal_tag = Tag("py3", "none", "any") @@ -144,6 +180,7 @@ def test_prefers_better_supported_tag(self): self.assertEqual(selection.name, "ptoas-0.57-cp312-cp312-manylinux_2_34_x86_64.whl") + @unittest.skipIf(Tag is None, "packaging is not available in this Python environment") def test_rejects_missing_compatible_wheel(self): release = {"tag_name": "nightly", "assets": []} with self.assertRaisesRegex(RuntimeError, "no compatible ptoas wheel"): diff --git a/tools/install_nightly_wheel.py b/tools/install_nightly_wheel.py index 19329a93f6..5ad3cf53b5 100644 --- a/tools/install_nightly_wheel.py +++ b/tools/install_nightly_wheel.py @@ -131,13 +131,29 @@ def parse_digest(value: object) -> str | None: return None -def select_wheel(release: object, package: str = "ptoas") -> WheelSelection: +def packaging_modules(): try: from packaging.tags import sys_tags from packaging.utils import canonicalize_name, parse_wheel_filename - except ImportError as error: + return sys_tags, canonicalize_name, parse_wheel_filename + except ImportError: + try: + from pip._vendor.packaging.tags import sys_tags + from pip._vendor.packaging.utils import canonicalize_name, parse_wheel_filename + return sys_tags, canonicalize_name, parse_wheel_filename + except ImportError as error: + raise RuntimeError( + "packaging support is required; use a Python environment with pip " + "or install it with 'python -m pip install packaging'" + ) from error + + +def select_wheel(release: object, package: str = "ptoas") -> WheelSelection: + try: + sys_tags, canonicalize_name, parse_wheel_filename = packaging_modules() + except RuntimeError as error: raise RuntimeError( - "the packaging module is required; install it with 'python -m pip install packaging'" + str(error) ) from error if not isinstance(release, dict) or not isinstance(release.get("assets"), list): From 5ebfd61054768d184b4194765f569ba21dfaa8aa Mon Sep 17 00:00:00 2001 From: jimmychou <47636600+jimmychou0@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:19:27 +0800 Subject: [PATCH 034/122] feat(vpto): make MX staging controls optional --- docs/isa/micro-isa/16-cube-matmul.md | 43 ++- include/PTO/IR/VPTOOps.td | 36 +- lib/PTO/IR/VPTO.cpp | 329 ++++++++++++++++-- lib/PTO/Transforms/VPTOExpandWrapperOps.cpp | 42 ++- .../user_guide/04-type-system-and-buffer.md | 2 +- .../docs/user_guide/07-data-movement-ops.md | 32 +- ptodsl/ptodsl/_ops.py | 110 +++--- ptodsl/tests/test_vector_cube_ops.py | 119 +++++-- .../cube/load_cbuf_to_mx_verify_invalid.pto | 61 ++++ .../cube/mte_l1_l0_mx_optional_operands.pto | 50 +++ 10 files changed, 673 insertions(+), 151 deletions(-) create mode 100644 test/lit/vpto/cube/mte_l1_l0_mx_optional_operands.pto diff --git a/docs/isa/micro-isa/16-cube-matmul.md b/docs/isa/micro-isa/16-cube-matmul.md index eb1730f99b..16105088dc 100644 --- a/docs/isa/micro-isa/16-cube-matmul.md +++ b/docs/isa/micro-isa/16-cube-matmul.md @@ -682,6 +682,39 @@ entry applies to one 32-element K group. - L1 source data is organized as 32B scale fragments in the same logical order as the associated data tile. +### MX Scale Load Operands + +`pto.mte_l1_l0a_mx` and `pto.mte_l1_l0b_mx` store the following operands as +independently optional fields, but a valid operation supplies exactly one +complete operand group: + +- **Shape-derived group:** `m`, `k`, `start_row`, `start_col` for L0A, or + `k`, `n`, `start_row`, `start_col` for L0B. PTOAS derives the MX traversal + from the matrix shape and element type. +- **Full MX group:** `x_start`, `y_start`, `x_step`, `y_step`, `src_stride`, + `dst_stride`. Use this group when the scale source has an explicit physical + layout that cannot be derived from the logical matrix shape. + +The two groups are mutually exclusive. A partial group is invalid. `x_start` +and `y_start` are MX scale-fragment grid coordinates; `x_step` and `y_step` +are grid traversal extents; `src_stride` and `dst_stride` are grid-row strides. +The full MX values are not logical element counts or byte offsets. + +The named field spelling is available when an incomplete or mixed operation +must be diagnosed. Complete groups print in the compact positional spelling +used by the examples below. + +```mlir +pto.mte_l1_l0a_mx %src, %dst, + x_start(%x_start), y_start(%y_start), x_step(%x_step), y_step(%y_step), + src_stride(%src_stride), dst_stride(%dst_stride) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 +``` + +`%src` must be in `l1` and address a 32-byte-aligned MX scale fragment. +`%dst` must be in the matching `l0a` or `l0b` space. It is always the real +staged L0 byte pointer and must be 16-byte aligned. + ### `pto.mte_l1_l0a_mx` - **syntax:** @@ -706,7 +739,10 @@ pto.mte_l1_l0a_mx %src, %dst, %m, %k, %start_row, %start_col **Constraints:** - `%src` must be in `l1`, `%dst` must be in `l0a`. -- `%src` and `%dst` must satisfy 32B MX scale-fragment alignment. +- `%src` must be 32-byte aligned to an MX scale fragment; `%dst` must be + 16-byte aligned for the MX destination address unit. +- Use either the complete shape-derived group or the complete full MX group + from [MX Scale Load Operands](#mx-scale-load-operands). **Example:** @@ -741,7 +777,10 @@ pto.mte_l1_l0b_mx %src, %dst, %k, %n, %start_row, %start_col **Constraints:** - `%src` must be in `l1`, `%dst` must be in `l0b`. -- `%src` and `%dst` must satisfy 32B MX scale-fragment alignment. +- `%src` must be 32-byte aligned to an MX scale fragment; `%dst` must be + 16-byte aligned for the MX destination address unit. +- Use either the complete shape-derived group or the complete full MX group + from [MX Scale Load Operands](#mx-scale-load-operands). **Example:** diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index ef914668ea..f59dbd8f55 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -3099,41 +3099,53 @@ def PTO_MteL1L0bOp : PTO_MteOp<"mte_l1_l0b", [ } def PTO_MteL1L0aMxOp : PTO_MteOp<"mte_l1_l0a_mx", [ + AttrSizedOperandSegments, DeclareOpInterfaceMethods ]> { let arguments = (ins PTO_BufferLikeType:$source, PTO_BufferLikeType:$destination, - Variadic:$controls + Optional:$m, + Optional:$k, + Optional:$start_row, + Optional:$start_col, + Optional:$x_start, + Optional:$y_start, + Optional:$x_step, + Optional:$y_step, + Optional:$src_stride, + Optional:$dst_stride ); let results = (outs); let hasVerifier = 1; - - let assemblyFormat = [{ - $source `,` $destination `,` $controls attr-dict `:` type($source) `,` - type($destination) `,` type($controls) - }]; + let hasCustomAssemblyFormat = 1; } def PTO_MteL1L0bMxOp : PTO_MteOp<"mte_l1_l0b_mx", [ + AttrSizedOperandSegments, DeclareOpInterfaceMethods ]> { let arguments = (ins PTO_BufferLikeType:$source, PTO_BufferLikeType:$destination, - Variadic:$controls + Optional:$k, + Optional:$n, + Optional:$start_row, + Optional:$start_col, + Optional:$x_start, + Optional:$y_start, + Optional:$x_step, + Optional:$y_step, + Optional:$src_stride, + Optional:$dst_stride ); let results = (outs); let hasVerifier = 1; - - let assemblyFormat = [{ - $source `,` $destination `,` $controls attr-dict `:` type($source) `,` - type($destination) `,` type($controls) - }]; + let hasCustomAssemblyFormat = 1; } def PTO_MteL0cL1Op : PTO_MteOp<"mte_l0c_l1", [ diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index 90334e0a3e..0f50310ac6 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -7796,8 +7796,210 @@ static LogicalResult verifyCubeBridgeLoadStart(OpTy op) { "start_row", op.getStartCol(), "start_col"); } -static LogicalResult verifyMxLoadControls(Operation *op, - OperandRange controls) { +template +static void setMxLoadOperandSegmentSizes(OperationState &result, + ArrayRef segmentSizes) { + auto &segments = result.getOrAddProperties() + .operandSegmentSizes; + llvm::copy(segmentSizes, segments.begin()); +} + +struct MxLoadAsmOperand { + OpAsmParser::UnresolvedOperand operand; + Type type; + bool present = false; +}; + +static std::optional +getMxLoadOperandIndex(StringRef keyword, ArrayRef shapeNames) { + for (auto [index, name] : llvm::enumerate(shapeNames)) + if (keyword == name) + return index; + + static constexpr StringRef kFullNames[] = { + "x_start", "y_start", "x_step", "y_step", "src_stride", "dst_stride"}; + for (auto [index, name] : llvm::enumerate(kFullNames)) + if (keyword == name) + return shapeNames.size() + index; + return std::nullopt; +} + +template +static ParseResult parseMteL1L0MxOp(OpAsmParser &parser, + OperationState &result, + ArrayRef shapeNames) { + OpAsmParser::UnresolvedOperand source; + OpAsmParser::UnresolvedOperand destination; + if (parser.parseOperand(source) || parser.parseComma() || + parser.parseOperand(destination)) + return failure(); + + SmallVector legacyOperands; + SmallVector namedOperands(10); + SmallVector namedOperandOrder; + bool usesNamedOperands = false; + + auto parseNamedOperand = [&](StringRef keyword) -> ParseResult { + std::optional index = getMxLoadOperandIndex(keyword, shapeNames); + if (!index) + return parser.emitError(parser.getCurrentLocation(), + "unknown MX load operand '") + << keyword << "'"; + if (namedOperands[*index].present) + return parser.emitError(parser.getCurrentLocation(), + "duplicate MX load operand '") + << keyword << "'"; + if (parser.parseLParen() || parser.parseOperand(namedOperands[*index].operand) || + parser.parseRParen()) + return failure(); + namedOperands[*index].present = true; + namedOperandOrder.push_back(*index); + return success(); + }; + + if (succeeded(parser.parseOptionalComma())) { + StringRef keyword; + if (succeeded(parser.parseOptionalKeyword(&keyword))) { + usesNamedOperands = true; + if (parseNamedOperand(keyword)) + return failure(); + while (succeeded(parser.parseOptionalComma())) { + if (parser.parseKeyword(&keyword) || parseNamedOperand(keyword)) + return failure(); + } + } else { + OpAsmParser::UnresolvedOperand operand; + if (parser.parseOperand(operand)) + return failure(); + legacyOperands.push_back(operand); + while (succeeded(parser.parseOptionalComma())) { + if (parser.parseOperand(operand)) + return failure(); + legacyOperands.push_back(operand); + } + } + } + + if (!usesNamedOperands && legacyOperands.size() != 4 && + legacyOperands.size() != 6) + return parser.emitError(parser.getCurrentLocation(), + "expects either four shape-derived or six full " + "positional MX operands"); + + if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) + return failure(); + + Type sourceType; + Type destinationType; + if (parser.parseType(sourceType) || parser.parseComma() || + parser.parseType(destinationType)) + return failure(); + + SmallVector legacyTypes; + if (usesNamedOperands) { + for (unsigned index : namedOperandOrder) { + Type type; + if (parser.parseComma() || parser.parseType(type)) + return failure(); + namedOperands[index].type = type; + } + } else { + for (size_t index = 0; index < legacyOperands.size(); ++index) { + Type type; + if (parser.parseComma() || parser.parseType(type)) + return failure(); + legacyTypes.push_back(type); + } + } + + SmallVector segmentSizes(12, 0); + segmentSizes[0] = 1; + segmentSizes[1] = 1; + + if (parser.resolveOperand(source, sourceType, result.operands) || + parser.resolveOperand(destination, destinationType, result.operands)) + return failure(); + + if (usesNamedOperands) { + for (unsigned index = 0; index < namedOperands.size(); ++index) { + if (!namedOperands[index].present) + continue; + segmentSizes[2 + index] = 1; + if (parser.resolveOperand(namedOperands[index].operand, + namedOperands[index].type, + result.operands)) + return failure(); + } + } else { + const unsigned base = legacyOperands.size() <= 4 ? 0 : 4; + for (unsigned index = 0; index < legacyOperands.size(); ++index) { + segmentSizes[2 + base + index] = 1; + if (parser.resolveOperand(legacyOperands[index], legacyTypes[index], + result.operands)) + return failure(); + } + } + setMxLoadOperandSegmentSizes(result, segmentSizes); + return success(); +} + +static void printMteL1L0MxOp(OpAsmPrinter &printer, Operation *operation, + Value source, Value destination, + ArrayRef shapeOperands, + ArrayRef shapeNames, + ArrayRef fullOperands) { + SmallVector fullNames = { + "x_start", "y_start", "x_step", "y_step", "src_stride", "dst_stride"}; + + const bool hasShape = llvm::any_of(shapeOperands, [](Value value) { + return static_cast(value); + }); + const bool hasFull = llvm::any_of(fullOperands, [](Value value) { + return static_cast(value); + }); + const bool isShapeForm = hasShape && !hasFull && + llvm::all_of(shapeOperands, [](Value value) { return static_cast(value); }); + const bool isFullForm = hasFull && !hasShape && + llvm::all_of(fullOperands, [](Value value) { return static_cast(value); }); + + printer << " " << source << ", " << destination; + SmallVector printedOperands; + if (isShapeForm) { + for (Value value : shapeOperands) { + printer << ", " << value; + printedOperands.push_back(value); + } + } else if (isFullForm) { + for (Value value : fullOperands) { + printer << ", " << value; + printedOperands.push_back(value); + } + } else { + for (auto [index, value] : llvm::enumerate(shapeOperands)) { + if (!value) + continue; + printer << ", " << shapeNames[index] << "(" << value << ")"; + printedOperands.push_back(value); + } + for (auto [index, value] : llvm::enumerate(fullOperands)) { + if (!value) + continue; + printer << ", " << fullNames[index] << "(" << value << ")"; + printedOperands.push_back(value); + } + } + + printer.printOptionalAttrDict(operation->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); + printer << " : " << source.getType() << ", " << destination.getType(); + for (Value value : printedOperands) + printer << ", " << value.getType(); +} + +static LogicalResult verifyMxLoadOperands(Operation *op, + ArrayRef shapeOperands, + ArrayRef shapeNames, + ArrayRef fullOperands) { auto checkNonNegativeConst = [&](Value value, StringRef name) -> LogicalResult { APInt intValue; @@ -7806,40 +8008,71 @@ static LogicalResult verifyMxLoadControls(Operation *op, return success(); }; - if (controls.size() == 4) - return verifyCubeBridgeLoadStart(op, controls[2], "start_row", - controls[3], "start_col"); - if (controls.size() == 6) { - if (failed(verifyCubeBridgeLoadStart(op, controls[0], "x_start", - controls[1], "y_start"))) - return failure(); - if (failed(checkNonNegativeConst(controls[2], "x_step")) || - failed(checkNonNegativeConst(controls[3], "y_step")) || - failed(checkNonNegativeConst(controls[4], "src_stride")) || - failed(checkNonNegativeConst(controls[5], "dst_stride"))) - return failure(); - return success(); + const bool hasShape = llvm::any_of(shapeOperands, [](Value value) { + return static_cast(value); + }); + const bool hasFull = llvm::any_of(fullOperands, [](Value value) { + return static_cast(value); + }); + if (hasShape && hasFull) + return op->emitOpError() + << "cannot mix shape-derived MX operands with full MX operands"; + if (!hasShape && !hasFull) + return op->emitOpError() + << "requires either all shape-derived MX operands or all full MX operands"; + + if (hasShape) { + for (auto [value, name] : llvm::zip(shapeOperands, shapeNames)) + if (!value) + return op->emitOpError() + << "shape-derived MX form requires " << name; + return verifyCubeBridgeLoadStart(op, shapeOperands[2], shapeNames[2], + shapeOperands[3], shapeNames[3]); } - return op->emitOpError() - << "requires either four shape-derived controls or six explicit " - "MX controls"; + + static constexpr StringRef kFullNames[] = { + "x_start", "y_start", "x_step", "y_step", "src_stride", "dst_stride"}; + for (auto [value, name] : llvm::zip(fullOperands, kFullNames)) + if (!value) + return op->emitOpError() << "full MX form requires " << name; + if (failed(verifyCubeBridgeLoadStart(op, fullOperands[0], "x_start", + fullOperands[1], "y_start"))) + return failure(); + if (failed(checkNonNegativeConst(fullOperands[2], "x_step")) || + failed(checkNonNegativeConst(fullOperands[3], "y_step")) || + failed(checkNonNegativeConst(fullOperands[4], "src_stride")) || + failed(checkNonNegativeConst(fullOperands[5], "dst_stride"))) + return failure(); + return success(); } -static LogicalResult verifyMxDestinationAlignment(Operation *op, - Value destination) { - constexpr int64_t kMxDestinationAddressUnitBytes = 16; - auto pointerCast = destination.getDefiningOp(); +static LogicalResult verifyMxPointerAlignment(Operation *op, Value pointer, + StringRef pointerName, + int64_t alignmentBytes) { + auto pointerCast = pointer.getDefiningOp(); if (!pointerCast || !isa(pointerCast.getInput().getType())) return success(); std::optional address = mlir::getConstantIntValue(pointerCast.getInput()); - if (!address || (*address % kMxDestinationAddressUnitBytes) == 0) + if (!address || (*address % alignmentBytes) == 0) return success(); return op->emitOpError() - << "statically known LOAD.MX destination address must be aligned to " - << kMxDestinationAddressUnitBytes << " bytes, got " << *address; + << "statically known LOAD.MX " << pointerName + << " address must be aligned to " << alignmentBytes << " bytes, got " + << *address; +} + +static LogicalResult verifyMxLoadAlignment(Operation *op, Value source, + Value destination) { + constexpr int64_t kMxSourceAlignmentBytes = 32; + constexpr int64_t kMxDestinationAddressUnitBytes = 16; + if (failed(verifyMxPointerAlignment(op, source, "source", + kMxSourceAlignmentBytes))) + return failure(); + return verifyMxPointerAlignment(op, destination, "destination", + kMxDestinationAddressUnitBytes); } LogicalResult MteL0cL1Op::verify() { @@ -7873,20 +8106,58 @@ LogicalResult MteL1L0bOp::verify() { return verifyCubeBridgeLoadStart(*this); } +ParseResult MteL1L0aMxOp::parse(OpAsmParser &parser, OperationState &result) { + static constexpr StringRef kShapeNames[] = { + "m", "k", "start_row", "start_col"}; + return parseMteL1L0MxOp(parser, result, kShapeNames); +} + +void MteL1L0aMxOp::print(OpAsmPrinter &printer) { + static constexpr StringRef kShapeNames[] = { + "m", "k", "start_row", "start_col"}; + printMteL1L0MxOp(printer, getOperation(), getSource(), getDestination(), + {getM(), getK(), getStartRow(), getStartCol()}, kShapeNames, + {getXStart(), getYStart(), getXStep(), getYStep(), + getSrcStride(), getDstStride()}); +} + LogicalResult MteL1L0aMxOp::verify() { if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::LEFT, "LEFT"))) return failure(); - if (failed(verifyMxLoadControls(getOperation(), getControls()))) + if (failed(verifyMxLoadOperands( + getOperation(), {getM(), getK(), getStartRow(), getStartCol()}, + {"m", "k", "start_row", "start_col"}, + {getXStart(), getYStart(), getXStep(), getYStep(), getSrcStride(), + getDstStride()}))) return failure(); - return verifyMxDestinationAlignment(getOperation(), getDestination()); + return verifyMxLoadAlignment(getOperation(), getSource(), getDestination()); +} + +ParseResult MteL1L0bMxOp::parse(OpAsmParser &parser, OperationState &result) { + static constexpr StringRef kShapeNames[] = { + "k", "n", "start_row", "start_col"}; + return parseMteL1L0MxOp(parser, result, kShapeNames); +} + +void MteL1L0bMxOp::print(OpAsmPrinter &printer) { + static constexpr StringRef kShapeNames[] = { + "k", "n", "start_row", "start_col"}; + printMteL1L0MxOp(printer, getOperation(), getSource(), getDestination(), + {getK(), getN(), getStartRow(), getStartCol()}, kShapeNames, + {getXStart(), getYStart(), getXStep(), getYStep(), + getSrcStride(), getDstStride()}); } LogicalResult MteL1L0bMxOp::verify() { if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::RIGHT, "RIGHT"))) return failure(); - if (failed(verifyMxLoadControls(getOperation(), getControls()))) + if (failed(verifyMxLoadOperands( + getOperation(), {getK(), getN(), getStartRow(), getStartCol()}, + {"k", "n", "start_row", "start_col"}, + {getXStart(), getYStart(), getXStep(), getYStep(), getSrcStride(), + getDstStride()}))) return failure(); - return verifyMxDestinationAlignment(getOperation(), getDestination()); + return verifyMxLoadAlignment(getOperation(), getSource(), getDestination()); } LogicalResult LoadCbufToCaMxOp::verify() { diff --git a/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp b/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp index 23f2a43269..2bb75efcb4 100644 --- a/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp +++ b/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp @@ -1570,18 +1570,21 @@ struct ExpandLeftLoadMxPattern : public OpRewritePattern { return rewriter.notifyMatchFailure( op, "failed to derive MX scale destination pointer"); - OperandRange controls = op.getControls(); - if (controls.size() != 4 && controls.size() != 6) - return rewriter.notifyMatchFailure( - op, "expected four shape-derived controls or six explicit MX controls"); LoadCbufToMxControl control; - if (controls.size() == 6) { - control = {controls[0], controls[1], controls[2], controls[3], - controls[4], controls[5]}; + if (op.getXStart()) { + if (!op.getYStart() || !op.getXStep() || !op.getYStep() || + !op.getSrcStride() || !op.getDstStride()) + return rewriter.notifyMatchFailure(op, + "expected complete full MX operands"); + control = {op.getXStart(), op.getYStart(), op.getXStep(), op.getYStep(), + op.getSrcStride(), op.getDstStride()}; } else { + if (!op.getM() || !op.getK() || !op.getStartRow() || !op.getStartCol()) + return rewriter.notifyMatchFailure( + op, "expected complete shape-derived MX operands"); FailureOr derived = deriveLoadCbufToCaMxControl( - loc, controls[0], controls[1], sourceType.getElementType(), - controls[2], controls[3], rewriter); + loc, op.getM(), op.getK(), sourceType.getElementType(), + op.getStartRow(), op.getStartCol(), rewriter); if (failed(derived)) return rewriter.notifyMatchFailure( op, "failed to derive load_cbuf_to_ca_mx control"); @@ -1616,18 +1619,21 @@ struct ExpandRightLoadMxPattern : public OpRewritePattern { return rewriter.notifyMatchFailure( op, "failed to derive MX scale destination pointer"); - OperandRange controls = op.getControls(); - if (controls.size() != 4 && controls.size() != 6) - return rewriter.notifyMatchFailure( - op, "expected four shape-derived controls or six explicit MX controls"); LoadCbufToMxControl control; - if (controls.size() == 6) { - control = {controls[0], controls[1], controls[2], controls[3], - controls[4], controls[5]}; + if (op.getXStart()) { + if (!op.getYStart() || !op.getXStep() || !op.getYStep() || + !op.getSrcStride() || !op.getDstStride()) + return rewriter.notifyMatchFailure(op, + "expected complete full MX operands"); + control = {op.getXStart(), op.getYStart(), op.getXStep(), op.getYStep(), + op.getSrcStride(), op.getDstStride()}; } else { + if (!op.getK() || !op.getN() || !op.getStartRow() || !op.getStartCol()) + return rewriter.notifyMatchFailure( + op, "expected complete shape-derived MX operands"); FailureOr derived = deriveLoadCbufToCbMxControl( - loc, controls[0], controls[1], sourceType.getElementType(), - controls[2], controls[3], rewriter); + loc, op.getK(), op.getN(), sourceType.getElementType(), + op.getStartRow(), op.getStartCol(), rewriter); if (failed(derived)) return rewriter.notifyMatchFailure( op, "failed to derive load_cbuf_to_cb_mx control"); diff --git a/ptodsl/docs/user_guide/04-type-system-and-buffer.md b/ptodsl/docs/user_guide/04-type-system-and-buffer.md index 39a9cf9460..b67c8dc9d0 100644 --- a/ptodsl/docs/user_guide/04-type-system-and-buffer.md +++ b/ptodsl/docs/user_guide/04-type-system-and-buffer.md @@ -63,7 +63,7 @@ lp_vreg_ty = pto.vreg_type(256, pto.f8e4m3) Constructing scalar eager values or host tensor ABI contracts with a low-precision type is **not supported** — `pto.f8e4m3(1.0)` and `pto.tensor_spec(rank=2, dtype=pto.f8e4m3)` will raise an error. -`pto.f8e8m0` is storage-only. Use it for MX scale storage, including an L1 source pointer passed to the explicit-control overload of `pto.mte_l1_l0a_mx` or `pto.mte_l1_l0b_mx`; it is not a scalar arithmetic type. +`pto.f8e8m0` is storage-only. Use it for MX scale storage, including an L1 source pointer passed with the full MX field group to `pto.mte_l1_l0a_mx` or `pto.mte_l1_l0b_mx`; it is not a scalar arithmetic type. ### Integer literal guidance diff --git a/ptodsl/docs/user_guide/07-data-movement-ops.md b/ptodsl/docs/user_guide/07-data-movement-ops.md index 76557fdf7d..bd809cefdf 100644 --- a/ptodsl/docs/user_guide/07-data-movement-ops.md +++ b/ptodsl/docs/user_guide/07-data-movement-ops.md @@ -979,25 +979,27 @@ Cube compute step; it does not issue those transfers itself. --- -#### `pto.mte_l1_l0a_mx(src: PtrType, dst: PtrType, m: int, k: int, *, start_row: int = 0, start_col: int = 0) -> None` -#### `pto.mte_l1_l0b_mx(src: PtrType, dst: PtrType, k: int, n: int, *, start_row: int = 0, start_col: int = 0) -> None` -#### `pto.mte_l1_l0a_mx(src: PtrType, dst: PtrType, *, x_start: int, y_start: int, x_step: int, y_step: int, src_stride: int, dst_stride: int) -> None` -#### `pto.mte_l1_l0b_mx(src: PtrType, dst: PtrType, *, x_start: int, y_start: int, x_step: int, y_step: int, src_stride: int, dst_stride: int) -> None` +#### `pto.mte_l1_l0a_mx(src: PtrType, dst: PtrType, m: int | None = None, k: int | None = None, *, start_row: int | None = None, start_col: int | None = None, x_start: int | None = None, y_start: int | None = None, x_step: int | None = None, y_step: int | None = None, src_stride: int | None = None, dst_stride: int | None = None) -> None` +#### `pto.mte_l1_l0b_mx(src: PtrType, dst: PtrType, k: int | None = None, n: int | None = None, *, start_row: int | None = None, start_col: int | None = None, x_start: int | None = None, y_start: int | None = None, x_step: int | None = None, y_step: int | None = None, src_stride: int | None = None, dst_stride: int | None = None) -> None` -**Description**: MX-mode variants of `mte_l1_l0a` and `mte_l1_l0b`. The shape-derived overloads preserve the existing behavior. The explicit-control overloads are for scale staging where traversal must not be inferred from matrix shape; they preserve every supplied control value through the existing PTO wrapper before it expands to the internal raw MX load. +**Description**: MX-mode variants of `mte_l1_l0a` and `mte_l1_l0b`. Every control field is independently optional, but a PTODSL call must provide exactly one complete group: either the shape-derived fields or the full MX fields. The shape-derived group preserves existing behavior. Use the full group when scale traversal cannot be inferred from matrix shape; every supplied value is preserved. | Parameter | Type | Description | |-----------|------|-------------| -| `src` | `PtrType` (L1) | MX scale source pointer. `pto.f8e8m0` is supported as a storage element type. | -| `dst` | `PtrType` (L0A or L0B) | Real staged L0 destination pointer. Pass the actual L0 address; do not divide or otherwise encode it in PTODSL. | -| `m`, `k` or `k`, `n` | `int` | Shape-derived overload dimensions. Do not combine them with explicit controls. | -| `start_row`, `start_col` | `int` | Shape-derived overload source offsets. Do not combine them with explicit controls. | -| `x_start` | `int` | MX load x start position. | -| `y_start` | `int` | MX load y start position. | -| `x_step` | `int` | MX load x traversal step. | -| `y_step` | `int` | MX load y traversal step. | -| `src_stride` | `int` | MX load source stride. | -| `dst_stride` | `int` | MX load destination stride. | +| `src` | `PtrType` (L1) | 32-byte-aligned MX scale source pointer. `pto.f8e8m0` is supported as a storage element type. | +| `dst` | `PtrType` (L0A or L0B) | 16-byte-aligned real staged L0 destination pointer. Pass the actual L0 address; do not divide or otherwise encode it in PTODSL. | +| `m`, `k` or `k`, `n` | `int | None` | Shape-derived dimensions. Supply both dimensions and do not combine them with full MX fields. | +| `start_row`, `start_col` | `int | None` | Shape-derived source offsets. They default to `0` when the shape-derived group is used. With full MX fields they must be omitted, except explicit literal `0` remains accepted for compatibility. | +| `x_start`, `y_start` | `int | None` | Start coordinates in the MX scale-fragment grid. | +| `x_step`, `y_step` | `int | None` | Traversal extents in MX scale-fragment grid units. | +| `src_stride` | `int | None` | Distance between source traversal rows in MX scale-fragment grid units. | +| `dst_stride` | `int | None` | Distance between destination traversal rows in MX scale-fragment grid units. | + +The full group is `x_start`, `y_start`, `x_step`, `y_step`, `src_stride`, and +`dst_stride`; all six values are required together. Its values are MX +scale-fragment grid coordinates or strides, not logical element counts or byte +offsets. A partial group or a mixture of shape-derived and full fields raises +`TypeError`. **Returns**: None (side-effect operation). diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 013f4936e4..a60d9ac9a1 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -5403,8 +5403,8 @@ def mte_l1_l0a_mx( m=None, k=None, *, - start_row=0, - start_col=0, + start_row=None, + start_col=None, x_start=None, y_start=None, x_step=None, @@ -5415,48 +5415,51 @@ def mte_l1_l0a_mx( """``pto.mte_l1_l0a_mx`` – MX cube-side LEFT staging. Use either the existing shape-derived ``m``/``k`` form or provide all six - explicit MX controls. Both forms trace through the L1-to-L0A MX wrapper. + full MX operands. Both forms trace through the L1-to-L0A MX wrapper. """ - controls = (x_start, y_start, x_step, y_step, src_stride, dst_stride) - has_explicit_controls = any(control is not None for control in controls) - if has_explicit_controls: + full_operands = (x_start, y_start, x_step, y_step, src_stride, dst_stride) + has_full_operands = any(operand is not None for operand in full_operands) + if has_full_operands: if m is not None or k is not None: raise TypeError( - "mte_l1_l0a_mx accepts either m/k or explicit MX controls, not both" + "mte_l1_l0a_mx accepts either m/k or full MX operands, not both" ) - if start_row != 0 or start_col != 0: + # The legacy API defaulted these shape-only fields to zero. Continue + # accepting explicit zero so existing full-MX callers remain valid. + if ((start_row is not None and start_row != 0) or + (start_col is not None and start_col != 0)): raise TypeError( - "mte_l1_l0a_mx start_row/start_col are unavailable with explicit MX controls" + "mte_l1_l0a_mx start_row/start_col are unavailable with full MX operands" ) - if any(control is None for control in controls): + if any(operand is None for operand in full_operands): raise TypeError( - "mte_l1_l0a_mx explicit MX controls require x_start, y_start, " + "mte_l1_l0a_mx full MX operands require x_start, y_start, " "x_step, y_step, src_stride, and dst_stride" ) _pto.MteL1L0aMxOp( unwrap_surface_value(source), unwrap_surface_value(destination), - [ - _coerce_i64(x_start, context="mte_l1_l0a_mx x_start"), - _coerce_i64(y_start, context="mte_l1_l0a_mx y_start"), - _coerce_i64(x_step, context="mte_l1_l0a_mx x_step"), - _coerce_i64(y_step, context="mte_l1_l0a_mx y_step"), - _coerce_i64(src_stride, context="mte_l1_l0a_mx src_stride"), - _coerce_i64(dst_stride, context="mte_l1_l0a_mx dst_stride"), - ], + x_start=_coerce_i64(x_start, context="mte_l1_l0a_mx x_start"), + y_start=_coerce_i64(y_start, context="mte_l1_l0a_mx y_start"), + x_step=_coerce_i64(x_step, context="mte_l1_l0a_mx x_step"), + y_step=_coerce_i64(y_step, context="mte_l1_l0a_mx y_step"), + src_stride=_coerce_i64(src_stride, context="mte_l1_l0a_mx src_stride"), + dst_stride=_coerce_i64(dst_stride, context="mte_l1_l0a_mx dst_stride"), ) return if m is None or k is None: - raise TypeError("mte_l1_l0a_mx requires m and k without explicit MX controls") + raise TypeError("mte_l1_l0a_mx requires m and k without full MX operands") + if start_row is None: + start_row = 0 + if start_col is None: + start_col = 0 _pto.MteL1L0aMxOp( unwrap_surface_value(source), unwrap_surface_value(destination), - [ - _coerce_i64(m, context="mte_l1_l0a_mx m"), - _coerce_i64(k, context="mte_l1_l0a_mx k"), - _coerce_i64(start_row, context="mte_l1_l0a_mx start_row"), - _coerce_i64(start_col, context="mte_l1_l0a_mx start_col"), - ], + m=_coerce_i64(m, context="mte_l1_l0a_mx m"), + k=_coerce_i64(k, context="mte_l1_l0a_mx k"), + start_row=_coerce_i64(start_row, context="mte_l1_l0a_mx start_row"), + start_col=_coerce_i64(start_col, context="mte_l1_l0a_mx start_col"), ) @@ -5467,8 +5470,8 @@ def mte_l1_l0b_mx( k=None, n=None, *, - start_row=0, - start_col=0, + start_row=None, + start_col=None, x_start=None, y_start=None, x_step=None, @@ -5479,48 +5482,51 @@ def mte_l1_l0b_mx( """``pto.mte_l1_l0b_mx`` – MX cube-side RIGHT staging. Use either the existing shape-derived ``k``/``n`` form or provide all six - explicit MX controls. Both forms trace through the L1-to-L0B MX wrapper. + full MX operands. Both forms trace through the L1-to-L0B MX wrapper. """ - controls = (x_start, y_start, x_step, y_step, src_stride, dst_stride) - has_explicit_controls = any(control is not None for control in controls) - if has_explicit_controls: + full_operands = (x_start, y_start, x_step, y_step, src_stride, dst_stride) + has_full_operands = any(operand is not None for operand in full_operands) + if has_full_operands: if k is not None or n is not None: raise TypeError( - "mte_l1_l0b_mx accepts either k/n or explicit MX controls, not both" + "mte_l1_l0b_mx accepts either k/n or full MX operands, not both" ) - if start_row != 0 or start_col != 0: + # The legacy API defaulted these shape-only fields to zero. Continue + # accepting explicit zero so existing full-MX callers remain valid. + if ((start_row is not None and start_row != 0) or + (start_col is not None and start_col != 0)): raise TypeError( - "mte_l1_l0b_mx start_row/start_col are unavailable with explicit MX controls" + "mte_l1_l0b_mx start_row/start_col are unavailable with full MX operands" ) - if any(control is None for control in controls): + if any(operand is None for operand in full_operands): raise TypeError( - "mte_l1_l0b_mx explicit MX controls require x_start, y_start, " + "mte_l1_l0b_mx full MX operands require x_start, y_start, " "x_step, y_step, src_stride, and dst_stride" ) _pto.MteL1L0bMxOp( unwrap_surface_value(source), unwrap_surface_value(destination), - [ - _coerce_i64(x_start, context="mte_l1_l0b_mx x_start"), - _coerce_i64(y_start, context="mte_l1_l0b_mx y_start"), - _coerce_i64(x_step, context="mte_l1_l0b_mx x_step"), - _coerce_i64(y_step, context="mte_l1_l0b_mx y_step"), - _coerce_i64(src_stride, context="mte_l1_l0b_mx src_stride"), - _coerce_i64(dst_stride, context="mte_l1_l0b_mx dst_stride"), - ], + x_start=_coerce_i64(x_start, context="mte_l1_l0b_mx x_start"), + y_start=_coerce_i64(y_start, context="mte_l1_l0b_mx y_start"), + x_step=_coerce_i64(x_step, context="mte_l1_l0b_mx x_step"), + y_step=_coerce_i64(y_step, context="mte_l1_l0b_mx y_step"), + src_stride=_coerce_i64(src_stride, context="mte_l1_l0b_mx src_stride"), + dst_stride=_coerce_i64(dst_stride, context="mte_l1_l0b_mx dst_stride"), ) return if k is None or n is None: - raise TypeError("mte_l1_l0b_mx requires k and n without explicit MX controls") + raise TypeError("mte_l1_l0b_mx requires k and n without full MX operands") + if start_row is None: + start_row = 0 + if start_col is None: + start_col = 0 _pto.MteL1L0bMxOp( unwrap_surface_value(source), unwrap_surface_value(destination), - [ - _coerce_i64(k, context="mte_l1_l0b_mx k"), - _coerce_i64(n, context="mte_l1_l0b_mx n"), - _coerce_i64(start_row, context="mte_l1_l0b_mx start_row"), - _coerce_i64(start_col, context="mte_l1_l0b_mx start_col"), - ], + k=_coerce_i64(k, context="mte_l1_l0b_mx k"), + n=_coerce_i64(n, context="mte_l1_l0b_mx n"), + start_row=_coerce_i64(start_row, context="mte_l1_l0b_mx start_row"), + start_col=_coerce_i64(start_col, context="mte_l1_l0b_mx start_col"), ) diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index d63699c96a..ab134cb99e 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -21,7 +21,7 @@ def _identity(value): class VectorCubeSurfaceTest(unittest.TestCase): - def test_mte_mx_explicit_controls_preserve_all_values(self): + def test_mte_mx_full_operands_preserve_all_values(self): source = object() destination = object() controls = { @@ -36,41 +36,99 @@ def test_mte_mx_explicit_controls_preserve_all_values(self): def coerce(value, *, context): return f"{context}:{value}" - expected = ( + with patch.object(_ops, "_require_explicit_mode"), \ + patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "_coerce_i64", side_effect=coerce), \ + patch.object(_ops._pto, "MteL1L0aMxOp") as load_ca: + pto.mte_l1_l0a_mx(source, destination, **controls) + load_ca.assert_called_once_with( source, destination, - "mte_l1_l0a_mx x_start:3", - "mte_l1_l0a_mx y_start:5", - "mte_l1_l0a_mx x_step:16", - "mte_l1_l0a_mx y_step:2", - "mte_l1_l0a_mx src_stride:8", - "mte_l1_l0a_mx dst_stride:2", + x_start="mte_l1_l0a_mx x_start:3", + y_start="mte_l1_l0a_mx y_start:5", + x_step="mte_l1_l0a_mx x_step:16", + y_step="mte_l1_l0a_mx y_step:2", + src_stride="mte_l1_l0a_mx src_stride:8", + dst_stride="mte_l1_l0a_mx dst_stride:2", ) + with patch.object(_ops, "_require_explicit_mode"), \ patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ patch.object(_ops, "_coerce_i64", side_effect=coerce), \ - patch.object(_ops._pto, "MteL1L0aMxOp") as load_ca: - pto.mte_l1_l0a_mx(source, destination, **controls) - load_ca.assert_called_once_with(source, destination, list(expected[2:])) + patch.object(_ops._pto, "MteL1L0bMxOp") as load_cb: + pto.mte_l1_l0b_mx(source, destination, **controls) + load_cb.assert_called_once_with( + source, + destination, + x_start="mte_l1_l0b_mx x_start:3", + y_start="mte_l1_l0b_mx y_start:5", + x_step="mte_l1_l0b_mx x_step:16", + y_step="mte_l1_l0b_mx y_step:2", + src_stride="mte_l1_l0b_mx src_stride:8", + dst_stride="mte_l1_l0b_mx dst_stride:2", + ) - expected = ( + def test_mte_mx_shape_operands_are_named(self): + source = object() + destination = object() + + def coerce(value, *, context): + return f"{context}:{value}" + + with patch.object(_ops, "_require_explicit_mode"), \ + patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "_coerce_i64", side_effect=coerce), \ + patch.object(_ops._pto, "MteL1L0aMxOp") as load_ca: + pto.mte_l1_l0a_mx( + source, destination, 128, 256, start_row=3, start_col=5 + ) + load_ca.assert_called_once_with( source, destination, - "mte_l1_l0b_mx x_start:3", - "mte_l1_l0b_mx y_start:5", - "mte_l1_l0b_mx x_step:16", - "mte_l1_l0b_mx y_step:2", - "mte_l1_l0b_mx src_stride:8", - "mte_l1_l0b_mx dst_stride:2", + m="mte_l1_l0a_mx m:128", + k="mte_l1_l0a_mx k:256", + start_row="mte_l1_l0a_mx start_row:3", + start_col="mte_l1_l0a_mx start_col:5", ) + with patch.object(_ops, "_require_explicit_mode"), \ patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ patch.object(_ops, "_coerce_i64", side_effect=coerce), \ patch.object(_ops._pto, "MteL1L0bMxOp") as load_cb: - pto.mte_l1_l0b_mx(source, destination, **controls) - load_cb.assert_called_once_with(source, destination, list(expected[2:])) + pto.mte_l1_l0b_mx( + source, destination, 256, 128, start_row=5, start_col=3 + ) + load_cb.assert_called_once_with( + source, + destination, + k="mte_l1_l0b_mx k:256", + n="mte_l1_l0b_mx n:128", + start_row="mte_l1_l0b_mx start_row:5", + start_col="mte_l1_l0b_mx start_col:3", + ) - def test_mte_mx_explicit_controls_require_a_complete_mode(self): + def test_mte_mx_shape_operands_default_starts_to_zero(self): + source = object() + destination = object() + + def coerce(value, *, context): + return f"{context}:{value}" + + with patch.object(_ops, "_require_explicit_mode"), \ + patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "_coerce_i64", side_effect=coerce), \ + patch.object(_ops._pto, "MteL1L0aMxOp") as load_ca: + pto.mte_l1_l0a_mx(source, destination, 128, 256) + load_ca.assert_called_once_with( + source, + destination, + m="mte_l1_l0a_mx m:128", + k="mte_l1_l0a_mx k:256", + start_row="mte_l1_l0a_mx start_row:0", + start_col="mte_l1_l0a_mx start_col:0", + ) + + def test_mte_mx_full_operands_require_a_complete_mode(self): source = object() destination = object() complete_controls = { @@ -85,8 +143,25 @@ def test_mte_mx_explicit_controls_require_a_complete_mode(self): with patch.object(_ops, "_require_explicit_mode"): with self.assertRaisesRegex(TypeError, "require x_start"): pto.mte_l1_l0a_mx(source, destination, x_start=3) - with self.assertRaisesRegex(TypeError, "either k/n or explicit"): + with self.assertRaisesRegex(TypeError, "either k/n or full"): pto.mte_l1_l0b_mx(source, destination, 128, 256, **complete_controls) + with patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "_coerce_i64", side_effect=lambda value, *, context: f"{context}:{value}"), \ + patch.object(_ops._pto, "MteL1L0aMxOp") as load_ca: + pto.mte_l1_l0a_mx( + source, destination, start_row=0, start_col=0, + **complete_controls + ) + load_ca.assert_called_once_with( + source, + destination, + x_start="mte_l1_l0a_mx x_start:3", + y_start="mte_l1_l0a_mx y_start:5", + x_step="mte_l1_l0a_mx x_step:16", + y_step="mte_l1_l0a_mx y_step:2", + src_stride="mte_l1_l0a_mx src_stride:8", + dst_stride="mte_l1_l0a_mx dst_stride:2", + ) self.assertFalse(hasattr(pto, "load_cbuf_to_ca_mx")) self.assertFalse(hasattr(pto, "load_cbuf_to_cb_mx")) diff --git a/test/lit/vpto/cube/load_cbuf_to_mx_verify_invalid.pto b/test/lit/vpto/cube/load_cbuf_to_mx_verify_invalid.pto index f19240e13b..4efc2107fc 100644 --- a/test/lit/vpto/cube/load_cbuf_to_mx_verify_invalid.pto +++ b/test/lit/vpto/cube/load_cbuf_to_mx_verify_invalid.pto @@ -9,11 +9,15 @@ // RUN: split-file %s %t // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_source.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-SOURCE // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/cb_destination.pto -o - 2>&1 | FileCheck %s --check-prefix=CB-DESTINATION +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_source_alignment.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-SOURCE-ALIGNMENT // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_alignment.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-ALIGNMENT // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_negative_x_step.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-NEGATIVE-X-STEP // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_negative_y_step.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-NEGATIVE-Y-STEP // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_negative_src_stride.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-NEGATIVE-SRC-STRIDE // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_negative_dst_stride.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-NEGATIVE-DST-STRIDE +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_partial_shape_fields.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-PARTIAL-SHAPE +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/cb_mixed_fields.pto -o - 2>&1 | FileCheck %s --check-prefix=CB-MIXED-FIELDS +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_five_positional_fields.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-FIVE-POSITIONAL //--- ca_source.pto module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { @@ -52,6 +56,19 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @ca_source_alignment() attributes {pto.kernel} { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %src = pto.castptr %c1 : i64 -> !pto.ptr + %dst = pto.castptr %c0 : i64 -> !pto.ptr + pto.mte_l1_l0a_mx %src, %dst, %c0, %c0, %c0, %c0 + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + return + } +} + //--- ca_negative_x_step.pto module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { func.func @ca_negative_x_step() attributes {pto.kernel} { @@ -112,10 +129,54 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @ca_partial_shape_fields() attributes {pto.kernel} { + %c0 = arith.constant 0 : i64 + %c16 = arith.constant 16 : i64 + %c64 = arith.constant 64 : i64 + %src = pto.castptr %c0 : i64 -> !pto.ptr + %dst = pto.castptr %c0 : i64 -> !pto.ptr + pto.mte_l1_l0a_mx %src, %dst, m(%c16), k(%c64), start_row(%c0) + : !pto.ptr, !pto.ptr, i64, i64, i64 + return + } +} + +//--- cb_mixed_fields.pto +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @cb_mixed_fields() attributes {pto.kernel} { + %c0 = arith.constant 0 : i64 + %c16 = arith.constant 16 : i64 + %c64 = arith.constant 64 : i64 + %src = pto.castptr %c0 : i64 -> !pto.ptr + %dst = pto.castptr %c0 : i64 -> !pto.ptr + pto.mte_l1_l0b_mx %src, %dst, k(%c64), n(%c16), start_row(%c0), start_col(%c0), x_start(%c0) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + return + } +} + +//--- ca_five_positional_fields.pto +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @ca_five_positional_fields() attributes {pto.kernel} { + %c0 = arith.constant 0 : i64 + %src = pto.castptr %c0 : i64 -> !pto.ptr + %dst = pto.castptr %c0 : i64 -> !pto.ptr + pto.mte_l1_l0a_mx %src, %dst, %c0, %c0, %c0, %c0, %c0 + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + return + } +} + // CA-SOURCE: 'pto.load_cbuf_to_ca_mx' op requires MAT source // CB-DESTINATION: 'pto.load_cbuf_to_cb_mx' op requires RIGHT destination +// CA-SOURCE-ALIGNMENT: 'pto.mte_l1_l0a_mx' op statically known LOAD.MX source address must be aligned to 32 bytes, got 1 // CA-ALIGNMENT: 'pto.mte_l1_l0a_mx' op statically known LOAD.MX destination address must be aligned to 16 bytes, got 1 // CA-NEGATIVE-X-STEP: 'pto.mte_l1_l0a_mx' op x_step must be non-negative // CA-NEGATIVE-Y-STEP: 'pto.mte_l1_l0a_mx' op y_step must be non-negative // CA-NEGATIVE-SRC-STRIDE: 'pto.mte_l1_l0a_mx' op src_stride must be non-negative // CA-NEGATIVE-DST-STRIDE: 'pto.mte_l1_l0a_mx' op dst_stride must be non-negative +// CA-PARTIAL-SHAPE: 'pto.mte_l1_l0a_mx' op shape-derived MX form requires start_col +// CB-MIXED-FIELDS: 'pto.mte_l1_l0b_mx' op cannot mix shape-derived MX operands with full MX operands +// CA-FIVE-POSITIONAL: expects either four shape-derived or six full positional MX operands diff --git a/test/lit/vpto/cube/mte_l1_l0_mx_optional_operands.pto b/test/lit/vpto/cube/mte_l1_l0_mx_optional_operands.pto new file mode 100644 index 0000000000..465bd4c4f6 --- /dev/null +++ b/test/lit/vpto/cube/mte_l1_l0_mx_optional_operands.pto @@ -0,0 +1,50 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-before=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @mx_load_optional_operands( + %a_src: !pto.ptr, %a_dst: !pto.ptr, + %b_src: !pto.ptr, %b_dst: !pto.ptr) + attributes {pto.kernel} { + %c0 = arith.constant 0 : i64 + %c2 = arith.constant 2 : i64 + %c3 = arith.constant 3 : i64 + %c5 = arith.constant 5 : i64 + %c8 = arith.constant 8 : i64 + %c11 = arith.constant 11 : i64 + %c13 = arith.constant 13 : i64 + %c16 = arith.constant 16 : i64 + %c17 = arith.constant 17 : i64 + %c128 = arith.constant 128 : i64 + + pto.mte_l1_l0a_mx %a_src, %a_dst, m(%c16), k(%c128), start_row(%c0), start_col(%c2) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.mte_l1_l0b_mx %b_src, %b_dst, dst_stride(%c13), x_step(%c17), y_start(%c5), src_stride(%c11), y_step(%c8), x_start(%c3) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + pto.mte_l1_l0a_mx %a_src, %a_dst, %c3, %c5, %c17, %c8, %c11, %c13 + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + return + } +} + +// CHECK-LABEL: func.func @mx_load_optional_operands( +// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : i64 +// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : i64 +// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : i64 +// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : i64 +// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : i64 +// CHECK-DAG: %[[C11:.*]] = arith.constant 11 : i64 +// CHECK-DAG: %[[C13:.*]] = arith.constant 13 : i64 +// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : i64 +// CHECK-DAG: %[[C17:.*]] = arith.constant 17 : i64 +// CHECK-DAG: %[[C128:.*]] = arith.constant 128 : i64 +// CHECK: pto.mte_l1_l0a_mx %{{.*}}, %{{.*}}, %[[C16]], %[[C128]], %[[C0]], %[[C2]] +// CHECK: pto.mte_l1_l0b_mx %{{.*}}, %{{.*}}, %[[C3]], %[[C5]], %[[C17]], %[[C8]], %[[C11]], %[[C13]] +// CHECK: pto.mte_l1_l0a_mx %{{.*}}, %{{.*}}, %[[C3]], %[[C5]], %[[C17]], %[[C8]], %[[C11]], %[[C13]] From 823ec9082742473dc8b6498e5b3c8a3945be4200 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Thu, 6 Aug 2026 20:37:00 +0800 Subject: [PATCH 035/122] fix(vmi): use broadcast load for single group slot --- docs/designs/vmi-layout-lowering-cases.md | 8 +++++ lib/PTO/Transforms/VMIToVPTO.cpp | 36 +++++++++++++++++++ .../vmi_compact_load_store_group_alias.pto | 21 +++++++++-- 3 files changed, 63 insertions(+), 2 deletions(-) diff --git a/docs/designs/vmi-layout-lowering-cases.md b/docs/designs/vmi-layout-lowering-cases.md index 4515f6e003..dd7223fafa 100644 --- a/docs/designs/vmi-layout-lowering-cases.md +++ b/docs/designs/vmi-layout-lowering-cases.md @@ -2383,6 +2383,14 @@ for g = 0..7: out[group_off + g] = rhs_base[rhs_off + g] ``` +The single-group unit-stride case is special. A `group_slot_load` with +`num_groups = 1` has only one semantic scalar slot, so it lowers through +`pto.vlds {dist = "BRC_B*"}` rather than `pto.vsldb`. `vsldb` requires its +source base operand to be 32B aligned, while the scalar effective address only +needs the natural alignment of its element type. The broadcast load preserves +the physical vreg shape; only lane 0 is semantically live for the one-slot +layout. + If `source_group_stride != 1`, this packed `slots = 8` layout requires a strided/gather group-slot load materializer. Until that support exists, `group_slot_load` with `slots = 8` and non-unit stride must diagnose instead of diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index bf5031d02b..d0a2288dc6 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -3338,6 +3338,13 @@ std::optional getPointStoreDistToken(Type elementType) { return (Twine("1PT_B") + Twine(elementBits)).str(); } +std::optional getScalarBroadcastLoadDistToken(Type elementType) { + unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); + if (elementBits != 8 && elementBits != 16 && elementBits != 32) + return std::nullopt; + return (Twine("BRC_B") + Twine(elementBits)).str(); +} + struct VPTOCmpMode { StringRef mode; std::optional signedness; @@ -6492,6 +6499,35 @@ static LogicalResult lowerGroupSlotLoadParts( if (!stride || *stride != 1) return rewriter.notifyMatchFailure( op, "slots=8 group_slot_load requires constant unit stride"); + + // A single logical group only needs one scalar source element. VSLDB is + // a 32B block load and requires its base operand to be 32B aligned, which + // is too strong for a valid element-aligned pointer such as base + k. + // VLD BRC loads that scalar from the effective element address and + // broadcasts it; only lane 0 is semantically live in this layout. + if (numGroups == 1) { + std::optional dist = + getScalarBroadcastLoadDistToken(resultVMIType.getElementType()); + if (!dist) + return rewriter.notifyMatchFailure( + op, "single-slot group_slot_load requires supported BRC load " + "element width"); + if (resultTypes.size() != 1) + return rewriter.notifyMatchFailure( + op, "single-slot group_slot_load arity mismatch"); + auto vregType = dyn_cast(resultTypes.front()); + if (!vregType) + return rewriter.notifyMatchFailure( + op, "single-slot group_slot_load result must be vreg"); + results.push_back(rewriter + .create(op->getLoc(), vregType, + /*updated_base=*/Type{}, source, + offset, + rewriter.getStringAttr(*dist)) + .getResult()); + return success(); + } + for (auto [chunk, resultType] : llvm::enumerate(resultTypes)) { auto vregType = dyn_cast(resultType); if (!vregType) diff --git a/test/lit/vmi_new/vmi_compact_load_store_group_alias.pto b/test/lit/vmi_new/vmi_compact_load_store_group_alias.pto index 34e016da89..641f4e2e05 100644 --- a/test/lit/vmi_new/vmi_compact_load_store_group_alias.pto +++ b/test/lit/vmi_new/vmi_compact_load_store_group_alias.pto @@ -19,6 +19,17 @@ module { return } + func.func @compact_1_after_addptr(%src: !pto.ptr, %dst: !pto.ptr, + %element: index) { + %zero = arith.constant 0 : index + %src_element = pto.addptr %src, %element : !pto.ptr -> !pto.ptr + %value = pto.vmi.vload %src_element[%zero] + : !pto.ptr -> !pto.vmi.vreg<1xf32> + pto.vmi.vstore %value, %dst[%element] + : !pto.vmi.vreg<1xf32>, !pto.ptr + return + } + func.func @compact_2(%src: !pto.ptr, %dst: !pto.ptr, %off: index) { %value = pto.vmi.vload %src[%off] @@ -96,8 +107,14 @@ module { // LEGACY-SAME: {num_groups = 1 : i64} // LOWER-LABEL: func.func @compact_1( -// LOWER: pto.pset_b32 "PAT_VL1" -// LOWER: pto.vsldb +// LOWER: pto.vlds {{.*}} {dist = "BRC_B32"} +// LOWER: pto.vsts +// LOWER-NOT: pto.vmi. + +// LOWER-LABEL: func.func @compact_1_after_addptr( +// LOWER: %[[SRC_ELEMENT:.*]] = pto.addptr %arg0, %arg2 +// LOWER: pto.vlds %[[SRC_ELEMENT]][%{{.*}}] {dist = "BRC_B32"} +// LOWER-NOT: pto.vsldb // LOWER: pto.vsts // LOWER-NOT: pto.vmi. From 22b0a4c2af52d883e2f87e373a893178e67d954e Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Fri, 7 Aug 2026 10:09:10 +0800 Subject: [PATCH 036/122] feat(vmi): support integer min and max --- lib/PTO/IR/VMI.cpp | 26 +++++++----- .../Transforms/VMILowerUnifiedToLegacy.cpp | 10 ++++- .../vmi_new/vmi_min_max_integer_invalid.pto | 37 ----------------- .../vmi_new/vmi_to_vpto_min_max_integer.pto | 40 +++++++++++++++++++ 4 files changed, 64 insertions(+), 49 deletions(-) delete mode 100644 test/lit/vmi_new/vmi_min_max_integer_invalid.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_min_max_integer.pto diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp index b6faeaa84d..c59a8946b7 100644 --- a/lib/PTO/IR/VMI.cpp +++ b/lib/PTO/IR/VMI.cpp @@ -2889,12 +2889,15 @@ LogicalResult VMIVminOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(lhsType.getElementType())) - return emitOpError("requires floating-point-like VMI element type"); - if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + Type elementType = lhsType.getElementType(); + if (!isVMIFloatLikeType(elementType) && !isVMIAnyI8I16I32Type(elementType)) + return emitOpError( + "requires floating-point-like or i8, i16, or i32 VMI element type"); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, + resultType))) return failure(); - if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), - resultType, getPmode()))) + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, + getPmode()))) return failure(); return success(); } @@ -2903,12 +2906,15 @@ LogicalResult VMIVmaxOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(lhsType.getElementType())) - return emitOpError("requires floating-point-like VMI element type"); - if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + Type elementType = lhsType.getElementType(); + if (!isVMIFloatLikeType(elementType) && !isVMIAnyI8I16I32Type(elementType)) + return emitOpError( + "requires floating-point-like or i8, i16, or i32 VMI element type"); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, + resultType))) return failure(); - if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), - resultType, getPmode()))) + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, + getPmode()))) return failure(); return success(); } diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp index 7a9949b176..fe3319f120 100644 --- a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp +++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp @@ -1458,18 +1458,24 @@ void VMILowerUnifiedToLegacyPass::runOnOperation() { } if (auto vop = dyn_cast(op)) { + Type elemType = getVMIElementType(vop.getResult()); auto createLegacy = [&](Location loc, Type ty, Value lhs, Value rhs) -> Value { - return builder.create(loc, ty, lhs, rhs).getResult(); + if (isFloatType(elemType)) + return builder.create(loc, ty, lhs, rhs).getResult(); + return builder.create(loc, ty, lhs, rhs).getResult(); }; (void)lowerBinaryIgnoringMask(vop, createLegacy); continue; } if (auto vop = dyn_cast(op)) { + Type elemType = getVMIElementType(vop.getResult()); auto createLegacy = [&](Location loc, Type ty, Value lhs, Value rhs) -> Value { - return builder.create(loc, ty, lhs, rhs).getResult(); + if (isFloatType(elemType)) + return builder.create(loc, ty, lhs, rhs).getResult(); + return builder.create(loc, ty, lhs, rhs).getResult(); }; (void)lowerBinaryIgnoringMask(vop, createLegacy); continue; diff --git a/test/lit/vmi_new/vmi_min_max_integer_invalid.pto b/test/lit/vmi_new/vmi_min_max_integer_invalid.pto deleted file mode 100644 index 4d84fe7fe5..0000000000 --- a/test/lit/vmi_new/vmi_min_max_integer_invalid.pto +++ /dev/null @@ -1,37 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -// RUN: not pto-test-opt %s -split-input-file 2>&1 | FileCheck %s - -module { - func.func @vmi_minf_integer_invalid( - %lhs: !pto.vmi.vreg<128xi32>, - %rhs: !pto.vmi.vreg<128xi32>) { - %min = pto.vmi.vmin %lhs, %rhs - : !pto.vmi.vreg<128xi32>, !pto.vmi.vreg<128xi32> - -> !pto.vmi.vreg<128xi32> - return - } -} - -// CHECK: 'pto.vmi.vmin' op requires floating-point-like VMI element type - -// ----- - -module { - func.func @vmi_maxf_integer_invalid( - %lhs: !pto.vmi.vreg<128xi32>, - %rhs: !pto.vmi.vreg<128xi32>) { - %max = pto.vmi.vmax %lhs, %rhs - : !pto.vmi.vreg<128xi32>, !pto.vmi.vreg<128xi32> - -> !pto.vmi.vreg<128xi32> - return - } -} - -// CHECK: 'pto.vmi.vmax' op requires floating-point-like VMI element type diff --git a/test/lit/vmi_new/vmi_to_vpto_min_max_integer.pto b/test/lit/vmi_new/vmi_to_vpto_min_max_integer.pto new file mode 100644 index 0000000000..382f1d76a1 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_min_max_integer.pto @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy | FileCheck %s --check-prefix=LEGACY +// RUN: pto-test-opt %s -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s --check-prefix=LOWER + +module { + func.func @vmi_to_vpto_min_max_integer( + %lhs: !pto.vmi.vreg<64xi32>, + %rhs: !pto.vmi.vreg<64xi32>) + -> (!pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32>) { + %min = pto.vmi.vmin %lhs, %rhs + : !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32> + -> !pto.vmi.vreg<64xi32> + %max = pto.vmi.vmax %lhs, %rhs + : !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32> + -> !pto.vmi.vreg<64xi32> + return %min, %max : !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32> + } +} + +// LEGACY-LABEL: func.func @vmi_to_vpto_min_max_integer( +// LEGACY: %[[MIN:.*]] = pto.vmi.mini +// LEGACY: %[[MAX:.*]] = pto.vmi.maxi +// LEGACY: return %[[MIN]], %[[MAX]] +// LEGACY-NOT: pto.vmi.vmin +// LEGACY-NOT: pto.vmi.vmax + +// LOWER-LABEL: func.func @vmi_to_vpto_min_max_integer( +// LOWER: %[[MIN:.*]] = pto.vmin +// LOWER: %[[MAX:.*]] = pto.vmax +// LOWER: return %[[MIN]], %[[MAX]] +// LOWER-NOT: pto.vmi. +// LOWER-NOT: !pto.vmi. +// LOWER-NOT: unrealized_conversion_cast From c148f6e6c7b71c13f1b952f4970a020e7c430c7b Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Tue, 4 Aug 2026 20:17:48 +0800 Subject: [PATCH 037/122] fix(ptodsl): normalize signed vdup scalar operands (#1102) --- ptodsl/ptodsl/_ops.py | 6 +++++ ptodsl/tests/test_issue_1102_vdup.py | 40 ++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+) create mode 100644 ptodsl/tests/test_issue_1102_vdup.py diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index a60d9ac9a1..cd5d700a9b 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -2098,6 +2098,12 @@ def vdup(input_value, mask, position=None): raise TypeError("vdup(scalar, mask, position=...) does not support position; position is only valid for vector input") raw_input = _coerce_vdup_scalar_input(input_value, mask, context="vdup(scalar, mask)") result_type = _infer_vdup_scalar_result_type(raw_input, mask, context="vdup(scalar, mask)") + result_element_type = _pto.VRegType(result_type).element_type + raw_input = coerce_scalar_to_type( + raw_input, + result_element_type, + context="vdup(scalar, mask)", + ) normalized_position = None return wrap_surface_value( _pto.VdupOp( diff --git a/ptodsl/tests/test_issue_1102_vdup.py b/ptodsl/tests/test_issue_1102_vdup.py new file mode 100644 index 0000000000..ca3d16acc9 --- /dev/null +++ b/ptodsl/tests/test_issue_1102_vdup.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software; you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +from ptodsl import pto + + +def _compile_vdup(name, dtype, mask_bits, value): + @pto.jit(name=name, kernel_kind="vector", target="a5", mode="explicit") + def kernel(): + mask = getattr(pto, f"pset_b{mask_bits}")(pto.MaskPattern.ALL) + pto.vdup(pto.const(value, dtype=dtype), mask) + + return kernel.compile().mlir_text() + + +def main(): + cases = ( + ("ui8", pto.ui8, 8, 0x80, "i8, !pto.mask -> !pto.vreg<256xi8>"), + ("si16", pto.si16, 16, 0x8000, "i16, !pto.mask -> !pto.vreg<128xi16>"), + ("ui16", pto.ui16, 16, 0x8000, "i16, !pto.mask -> !pto.vreg<128xi16>"), + ("ui32", pto.ui32, 32, 0x80000000, "i32, !pto.mask -> !pto.vreg<64xi32>"), + ) + for name, dtype, mask_bits, value, expected_type in cases: + text = _compile_vdup(f"issue_1102_vdup_{name}", dtype, mask_bits, value) + if expected_type not in text: + raise AssertionError( + f"{name} vdup should normalize its scalar operand to the signless result element type; " + f"expected {expected_type!r} in:\n{text}" + ) + print("issue_1102_vdup: PASS") + + +if __name__ == "__main__": + main() From f7626e802fd84b69afd736e736bd56f3f4f4c97a Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Tue, 4 Aug 2026 20:25:24 +0800 Subject: [PATCH 038/122] style: fix issue 1102 test license header --- ptodsl/tests/test_issue_1102_vdup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ptodsl/tests/test_issue_1102_vdup.py b/ptodsl/tests/test_issue_1102_vdup.py index ca3d16acc9..ec9cf6d4f6 100644 --- a/ptodsl/tests/test_issue_1102_vdup.py +++ b/ptodsl/tests/test_issue_1102_vdup.py @@ -1,9 +1,9 @@ #!/usr/bin/env python3 # Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software; you can redistribute it and/or modify it under the terms and conditions of +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of # CANN Open Software License Agreement Version 2.0 (the "License"). # Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. From e1c4fd9af0f9d28d2bcbf75792c6f90738e3913b Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Wed, 5 Aug 2026 10:29:12 +0800 Subject: [PATCH 039/122] fix(ptodsl): preserve vdup integer signedness --- ptodsl/ptodsl/_ops.py | 6 ++---- ptodsl/tests/test_issue_1102_vdup.py | 29 +++++++++++++++++++++++----- 2 files changed, 26 insertions(+), 9 deletions(-) diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index cd5d700a9b..914c39695b 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -63,7 +63,6 @@ _materialize_integer_literal, _normalize_address_space, _resolve, - _strip_integer_signedness, mask_type, part_tensor_view_type, part_tensor_view_type_from_dims, @@ -2040,13 +2039,12 @@ def _infer_vdup_scalar_result_type(input_value, mask_value, *, context: str): scalar_type = scalar_raw.type mask_bits = _mask_granularity_bits(mask_value, context=context) if IntegerType.isinstance(scalar_type): - scalar_type = _strip_integer_signedness(scalar_raw) - scalar_width = IntegerType(scalar_type.type).width + scalar_width = IntegerType(scalar_type).width if scalar_width != mask_bits: raise TypeError( f"{context} expects scalar input width {scalar_width} to match mask granularity b{mask_bits}" ) - element_type = scalar_type.type + element_type = scalar_type elif F16Type.isinstance(scalar_type) or BF16Type.isinstance(scalar_type): if mask_bits != 16: raise TypeError(f"{context} expects f16/bf16 scalar input to pair with mask_b16, got mask_b{mask_bits}") diff --git a/ptodsl/tests/test_issue_1102_vdup.py b/ptodsl/tests/test_issue_1102_vdup.py index ec9cf6d4f6..83d8b749e8 100644 --- a/ptodsl/tests/test_issue_1102_vdup.py +++ b/ptodsl/tests/test_issue_1102_vdup.py @@ -19,20 +19,39 @@ def kernel(): return kernel.compile().mlir_text() +def _compile_vdup_consumer(): + @pto.jit(name="issue_1102_vdup_vor", kernel_kind="vector", target="a5", mode="explicit") + def kernel(): + base = pto.const(0, dtype=pto.ui64) + source = pto.castptr(base, pto.ptr(pto.ui16, "ub")) + mask = pto.pset_b16(pto.MaskPattern.ALL) + loaded = pto.vlds(source, pto.const(0)) + duplicated = pto.vdup(pto.ui16(1), mask) + pto.vor(loaded, duplicated, mask) + + return kernel.compile().mlir_text() + + def main(): cases = ( - ("ui8", pto.ui8, 8, 0x80, "i8, !pto.mask -> !pto.vreg<256xi8>"), - ("si16", pto.si16, 16, 0x8000, "i16, !pto.mask -> !pto.vreg<128xi16>"), - ("ui16", pto.ui16, 16, 0x8000, "i16, !pto.mask -> !pto.vreg<128xi16>"), - ("ui32", pto.ui32, 32, 0x80000000, "i32, !pto.mask -> !pto.vreg<64xi32>"), + ("i8", pto.i8, 8, 1, "i8, !pto.mask -> !pto.vreg<256xi8>"), + ("si8", pto.si8, 8, 1, "si8, !pto.mask -> !pto.vreg<256xsi8>"), + ("ui8", pto.ui8, 8, 1, "ui8, !pto.mask -> !pto.vreg<256xui8>"), + ("i16", pto.i16, 16, 1, "i16, !pto.mask -> !pto.vreg<128xi16>"), + ("si16", pto.si16, 16, 1, "si16, !pto.mask -> !pto.vreg<128xsi16>"), + ("ui16", pto.ui16, 16, 1, "ui16, !pto.mask -> !pto.vreg<128xui16>"), + ("i32", pto.i32, 32, 1, "i32, !pto.mask -> !pto.vreg<64xi32>"), + ("si32", pto.si32, 32, 1, "si32, !pto.mask -> !pto.vreg<64xsi32>"), + ("ui32", pto.ui32, 32, 1, "ui32, !pto.mask -> !pto.vreg<64xui32>"), ) for name, dtype, mask_bits, value, expected_type in cases: text = _compile_vdup(f"issue_1102_vdup_{name}", dtype, mask_bits, value) if expected_type not in text: raise AssertionError( - f"{name} vdup should normalize its scalar operand to the signless result element type; " + f"{name} vdup should preserve the scalar integer signedness in its result vector; " f"expected {expected_type!r} in:\n{text}" ) + _compile_vdup_consumer() print("issue_1102_vdup: PASS") From 305a2064c924e659e75f516700d54353b8f879a2 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 7 Aug 2026 14:40:23 +0800 Subject: [PATCH 040/122] fix(emitc): forward compatible GlobalTensor casts for MGATHER/MSCATTER The PTOToEmitCTypeConverter maps tensor_view / partition_tensor_view to a GlobalTensor opaque type with fully-dynamic Stride<-1,...> template params, because strides are not carried on the MLIR type. The static partition_view pattern, however, materializes a GlobalTensor with concrete static strides. Dialect conversion bridges the two with an unrealized_conversion_cast that the cleanup lowered to an emitc.cast, i.e. an invalid C-style cast between two GlobalTensor instantiations that have no converting constructor. The generated C++ then failed to compile for MGATHER / MSCATTER (issue #1165, a v0.55 regression exposed once these ops became tensor-view-native). Recognize such structurally-compatible GlobalTensor-to-GlobalTensor bridges in the cast cleanup (identical except static-vs-dynamic Shape/Stride template params) and forward the value instead of emitting a cast, so the static-stride GlobalTensor flows directly into the templated MGATHER / MSCATTER call. Add a NOCAST FileCheck pass to the mgather/mscatter lit test guarding against re-emission of the C-style GlobalTensor cast. --- lib/PTO/Transforms/PTOToEmitC.cpp | 65 +++++++++++++++++++ .../pto/mgather_mscatter_base_a3_emitc.pto | 8 +++ 2 files changed, 73 insertions(+) diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 271f9e50a8..b88258ba30 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -13669,6 +13669,61 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, populateBranchOpInterfaceTypeConversionPattern(patterns, typeConverter); } +// A cast between two `GlobalTensor<...>` opaque C++ types that are identical +// except that one side carries concrete Shape/Stride template values while the +// other uses the fully-dynamic `-1` placeholders (as produced by the +// PTOToEmitCTypeConverter, which cannot recover strides from the stride-less +// tensor_view type) has no valid C++ converting constructor. The values are +// interchangeable at every templated backend call site (e.g. MGATHER/MSCATTER), +// so such a bridge must forward the value rather than lower to an invalid +// C-style `emitc.cast`. +static bool areRefinableGlobalTensorTypes(Type a, Type b) { + auto oa = dyn_cast(a); + auto ob = dyn_cast(b); + if (!oa || !ob) + return false; + StringRef sa = oa.getValue(); + StringRef sb = ob.getValue(); + if (!sa.contains("GlobalTensor<") || !sb.contains("GlobalTensor<")) + return false; + + SmallVector shapeA, shapeB, strideA, strideB; + if (!parseIntegerTemplateList(sa, "Shape<", shapeA) || + !parseIntegerTemplateList(sb, "Shape<", shapeB) || + !parseIntegerTemplateList(sa, "Stride<", strideA) || + !parseIntegerTemplateList(sb, "Stride<", strideB)) + return false; + + // Element-wise compatible if equal or one side is the `-1` wildcard. + auto listsRefinable = [](ArrayRef x, ArrayRef y) { + if (x.size() != y.size()) + return false; + for (auto [u, v] : llvm::zip(x, y)) + if (u != v && u != -1 && v != -1) + return false; + return true; + }; + if (!listsRefinable(shapeA, shapeB) || !listsRefinable(strideA, strideB)) + return false; + + // Everything outside the Shape<...>/Stride<...> lists (element type, layout, + // overall structure) must match exactly. + auto blank = [](StringRef s, StringRef marker) -> std::string { + std::string out = s.str(); + size_t pos = out.find(marker.str()); + if (pos == std::string::npos) + return out; + size_t start = pos + marker.size(); + size_t end = out.find('>', start); + if (end == std::string::npos) + return out; + out.erase(start, end - start); + return out; + }; + return blank(blank(sa, "Shape<"), "Stride<") == + blank(blank(sb, "Shape<"), "Stride<"); +} + //===----------------------------------------------------------------------===// // Pass //===----------------------------------------------------------------------===// @@ -14091,6 +14146,16 @@ static AICORE inline void PTOAS__DCCI_SINGLE_CACHE_LINE(Ptr ptr) { return; } + // A static-stride `GlobalTensor` produced by the partition_view static + // pattern and the fully-dynamic-stride `GlobalTensor` demanded by the + // type converter have no C++ converting constructor. Forward the value + // instead of emitting an invalid C-style cast. + if (areRefinableGlobalTensorTypes(inTy, outTy)) { + output.replaceAllUsesWith(input); + castsToErase.push_back(cast); + return; + } + if (emitc::isSupportedEmitCType(inTy) && emitc::isSupportedEmitCType(outTy)) { OpBuilder builder(cast); auto c = builder.create(cast.getLoc(), outTy, input); diff --git a/test/lit/pto/mgather_mscatter_base_a3_emitc.pto b/test/lit/pto/mgather_mscatter_base_a3_emitc.pto index e407a8e9db..0f7ef69b99 100644 --- a/test/lit/pto/mgather_mscatter_base_a3_emitc.pto +++ b/test/lit/pto/mgather_mscatter_base_a3_emitc.pto @@ -6,6 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s --check-prefix=NOCAST module attributes {pto.target_arch = "a3"} { func.func @mgather_emitc_base(%src: !pto.ptr) @@ -67,3 +68,10 @@ module attributes {pto.target_arch = "a3"} { // CHECK-LABEL: AICORE void mscatter_emitc_base( // CHECK: MSCATTER({{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}); + +// Regression guard for issue #1165: the partition-view GlobalTensor must flow +// into MGATHER/MSCATTER directly. No static-stride -> dynamic-stride +// GlobalTensor C-style cast (`(GlobalTensor<...>)v`) may be emitted, since +// there is no C++ converting constructor between differing Stride<...> template +// instantiations and the generated C++ would fail to compile. +// NOCAST-NOT: (GlobalTensor< From ddd6b415559f8a1d2e02931b2efe0e96623aa4da Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 7 Aug 2026 15:15:12 +0800 Subject: [PATCH 041/122] fix(emitc): keep GlobalTensor cast when refined value feeds a return Forwarding the static-stride GlobalTensor into a func/emitc return breaks verification because the enclosing function's result type is fixed to the dynamic-stride form. Only forward when every consumer accepts a more-specific template instantiation (e.g. MGATHER/MSCATTER); otherwise fall through to the emitc.cast branch. Fixes a regression in issue31_partition_view_parser_compat. --- lib/PTO/Transforms/PTOToEmitC.cpp | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index b88258ba30..d26ca6a7f4 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -14148,12 +14148,22 @@ static AICORE inline void PTOAS__DCCI_SINGLE_CACHE_LINE(Ptr ptr) { // A static-stride `GlobalTensor` produced by the partition_view static // pattern and the fully-dynamic-stride `GlobalTensor` demanded by the - // type converter have no C++ converting constructor. Forward the value - // instead of emitting an invalid C-style cast. + // type converter have no C++ converting constructor. Forwarding the + // refined (static) value is only safe when every consumer accepts a + // more-specific GlobalTensor template instantiation (e.g. MGATHER / + // MSCATTER, which are C++ templates). A `return` must match the enclosing + // function's fixed (dynamic) result type exactly, so forwarding there + // would break verification -- fall through to the emitc.cast branch in + // that case. if (areRefinableGlobalTensorTypes(inTy, outTy)) { - output.replaceAllUsesWith(input); - castsToErase.push_back(cast); - return; + bool feedsReturn = llvm::any_of(output.getUsers(), [](Operation *user) { + return isa(user); + }); + if (!feedsReturn) { + output.replaceAllUsesWith(input); + castsToErase.push_back(cast); + return; + } } if (emitc::isSupportedEmitCType(inTy) && emitc::isSupportedEmitCType(outTy)) { From c0a57210c10fdcd788a1b05d5334bbab29a3553e Mon Sep 17 00:00:00 2001 From: likai00 Date: Fri, 7 Aug 2026 16:26:50 +0800 Subject: [PATCH 042/122] sync docs --- docs/isa/vmi-isa/01-load-store.md | 3 +++ docs/isa/vmi-isa/02-index-gen.md | 1 + docs/isa/vmi-isa/03-eltwise-compute.md | 34 ++++++++++++++++++++------ docs/isa/vmi-isa/04-broadcast.md | 1 + docs/isa/vmi-isa/05-reduce.md | 2 ++ docs/isa/vmi-isa/06-convert.md | 2 ++ docs/isa/vmi-isa/07-sfu.md | 3 +++ docs/isa/vmi-isa/08-predicate-ops.md | 33 ++++++++++++++++++------- docs/isa/vmi-isa/09-data-rearrange.md | 2 ++ 9 files changed, 64 insertions(+), 17 deletions(-) diff --git a/docs/isa/vmi-isa/01-load-store.md b/docs/isa/vmi-isa/01-load-store.md index e53fc2d7da..a52fe12122 100644 --- a/docs/isa/vmi-isa/01-load-store.md +++ b/docs/isa/vmi-isa/01-load-store.md @@ -6,6 +6,7 @@ > the access pattern**, defaulting to `continuous` (contiguous); the optional > modes are `unpack` (widening unpack) and `brc` (broadcast). + --- ## `pto.vmi.vload` @@ -142,6 +143,7 @@ declaring the memory access pattern. Default is `"continuous"`. block-stride mode and invalid otherwise. `vload` has no mask operand in any mode (A5 loads are unpredicated). + --- ## `pto.vmi.vstore` @@ -238,6 +240,7 @@ declaring the memory access pattern. Default is `"continuous"`. // → block-strided store (block=8), governed by mask ``` + --- ## `pto.vmi.vsstb` diff --git a/docs/isa/vmi-isa/02-index-gen.md b/docs/isa/vmi-isa/02-index-gen.md index 478ba78982..55cdec3f1d 100644 --- a/docs/isa/vmi-isa/02-index-gen.md +++ b/docs/isa/vmi-isa/02-index-gen.md @@ -5,6 +5,7 @@ > Index materialization. Produces an index vector; the single physical reg > backing is replicate-read until a Category B/C edge needs the expanded form. + --- ## `pto.vmi.vci` diff --git a/docs/isa/vmi-isa/03-eltwise-compute.md b/docs/isa/vmi-isa/03-eltwise-compute.md index c7a5e72972..2f6dff4278 100644 --- a/docs/isa/vmi-isa/03-eltwise-compute.md +++ b/docs/isa/vmi-isa/03-eltwise-compute.md @@ -9,6 +9,7 @@ > expanded to `K` copies). Under the `K ≤ 4` core profile these fan out as > fully-unrolled straight-line code. + --- ## 3.1 Binary Arithmetic @@ -107,6 +108,7 @@ ``` `#mi = K`, `dep = 1`. + --- ## 3.2 Unary Arithmetic & Activation @@ -195,6 +197,7 @@ ``` `#mi = K`, `dep = 1`. + --- ## 3.3 Bitwise Ops @@ -202,9 +205,10 @@ ### `pto.vmi.vand` / `pto.vmi.vor` / `pto.vmi.vxor` - **semantics:** Elementwise bitwise AND / OR / XOR. Operands and result are - vregs by default; will also support mask-typed operands, performing a per-lane - predicate boolean op and yielding a mask (the data operands themselves are - masks, distinct from the governing `mask`). + vregs by default. These ops also accept mask-typed operands, performing a + per-lane predicate boolean op and yielding a mask. When the operands are + masks (predicate type), no governing `mask` operand may be given — a mask + operand would be ambiguous with the predicate data operands themselves. ```c for (int i = 0; i < L; i++) @@ -213,9 +217,14 @@ - **syntax:** ```mlir + // vreg operands (optional governing mask) %r = pto.vmi.vand %lhs, %rhs, %mask : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + + // mask operands (no governing mask) + %r = pto.vmi.vand %lhs, %rhs : !pto.vmi.mask, !pto.vmi.mask -> !pto.vmi.mask + %r = pto.vmi.vxor %lhs, %rhs : !pto.vmi.mask, !pto.vmi.mask -> !pto.vmi.mask ``` -- **datatypes:** `i8`–`i32` (integer bitwise) +- **datatypes:** `i8`–`i32` (integer bitwise); `pred` (per-lane boolean op) - **lowering to `pto.mi`:** ``` K × pto.vand / pto.vor / pto.vxor @@ -225,9 +234,10 @@ ### `pto.vmi.vnot` - **semantics:** Elementwise bitwise NOT. Operand and result are vregs by - default; will also support a mask-typed operand, performing a per-lane predicate - complement and yielding a mask (the data operand itself is a mask, distinct - from the governing `mask`). + default. This op also accepts a mask-typed operand, performing a per-lane + predicate complement and yielding a mask. When the operand is a mask + (predicate type), no governing `mask` operand may be given — a mask operand + would be ambiguous with the predicate data operand itself. ```c for (int i = 0; i < L; i++) @@ -236,15 +246,20 @@ - **syntax:** ```mlir + // vreg operand (optional governing mask) %r = pto.vmi.vnot %src, %mask : !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + + // mask operand (no governing mask) + %r = pto.vmi.vnot %src : !pto.vmi.mask -> !pto.vmi.mask ``` -- **datatypes:** `i8`–`i32` +- **datatypes:** `i8`–`i32`; `pred` (predicate complement) - **lowering to `pto.mi`:** ``` K × pto.vnot ``` `#mi = K`, `dep = 1`. + --- ## 3.4 Shift Ops @@ -271,6 +286,7 @@ ``` `#mi = K`, `dep = 1`. + --- ## 3.5 Vec-Scalar Ops @@ -343,6 +359,7 @@ scalar type must match the vector element type. ``` `#mi = K`, `dep = 1`. + --- ## 3.6 Compare & Select @@ -565,6 +582,7 @@ scalar type must match the vector element type. : !pto.vmi.vreg<128×f16>, !pto.vmi.vreg<128×i16> -> !pto.vmi.vreg<128×f16> ``` + --- ## 3.7 Carry / Borrow Ops (Not Provided) diff --git a/docs/isa/vmi-isa/04-broadcast.md b/docs/isa/vmi-isa/04-broadcast.md index 53e51f07dc..ec74277e5c 100644 --- a/docs/isa/vmi-isa/04-broadcast.md +++ b/docs/isa/vmi-isa/04-broadcast.md @@ -8,6 +8,7 @@ > (per-group scalar fan-back) has no single native instruction and is a > cost-model decision. + --- ## `pto.vmi.vbrc` diff --git a/docs/isa/vmi-isa/05-reduce.md b/docs/isa/vmi-isa/05-reduce.md index fd6e78d5d7..4d37058bdf 100644 --- a/docs/isa/vmi-isa/05-reduce.md +++ b/docs/isa/vmi-isa/05-reduce.md @@ -8,6 +8,7 @@ > `vcadd` treats inactive as 0; `vcmax`/`vcmin` treat inactive as `-∞`/`+∞` > (fp) or type min/max (int). + --- ## `pto.vmi.vcadd` @@ -79,6 +80,7 @@ : !pto.vmi.vreg<256×f16>, !pto.vmi.mask<256> -> !pto.vmi.vreg<8×f16> ``` + --- ## `pto.vmi.vcmax` / `pto.vmi.vcmin` diff --git a/docs/isa/vmi-isa/06-convert.md b/docs/isa/vmi-isa/06-convert.md index 192ef6bc15..edd8fcc4f9 100644 --- a/docs/isa/vmi-isa/06-convert.md +++ b/docs/isa/vmi-isa/06-convert.md @@ -8,6 +8,7 @@ > distribution. The author never spells `EVEN`/`ODD`, `P0`–`P3`, `PK`/`UNPK`, > or `VL/2` addresses. + --- ## `pto.vmi.vcvt` @@ -120,6 +121,7 @@ roundtrip; the 1↔4 lane spread rides data load/store distribution (`UNPK_B*`/`PK4_B32`) or a `vselr` byte-gather. + --- ## `pto.vmi.vinterpret_cast` diff --git a/docs/isa/vmi-isa/07-sfu.md b/docs/isa/vmi-isa/07-sfu.md index 815e008b01..3f84dcd501 100644 --- a/docs/isa/vmi-isa/07-sfu.md +++ b/docs/isa/vmi-isa/07-sfu.md @@ -9,6 +9,7 @@ > ops (including `vmull`, whose 64-bit product is split into a pair of `i32` > results at the VMI surface) are Category A `vreg→vreg`. + --- ## 7.1 Fused Arithmetic @@ -249,6 +250,7 @@ !pto.vmi.mask<64> -> !pto.vmi.vreg<64×f32> ``` + --- ## 7.2 Histogram @@ -401,6 +403,7 @@ -> !pto.vmi.vreg<256×i16> ``` + --- ## 7.3 Gather / Scatter diff --git a/docs/isa/vmi-isa/08-predicate-ops.md b/docs/isa/vmi-isa/08-predicate-ops.md index 1cea43a2f6..d2f5f13433 100644 --- a/docs/isa/vmi-isa/08-predicate-ops.md +++ b/docs/isa/vmi-isa/08-predicate-ops.md @@ -20,6 +20,7 @@ %next = arith.subi %rem, %act // rem - min(rem, L) ``` + --- ## `pto.vmi.create_mask` @@ -58,6 +59,7 @@ %tail = pto.vmi.create_mask %c64 : index -> !pto.vmi.mask<128×b32> ``` + --- ## `pto.vmi.create_group_mask` @@ -113,12 +115,25 @@ lane count. A backend may impose a narrower materialization limit separately. -> **Mask Boolean Ops (`vand` / `vor` / `vxor` / `vnot` on masks):** -> -> There is **no dedicated predicate-logic op** (e.g. `pand`/`por`/`pxor`/`pnot`). -> Mask (predicate) boolean operations are **not yet supported**, but are planned. -> The planned approach is to **reuse the elementwise bitwise ops** `pto.vmi.vand` / -> `vor` / `vxor` / `vnot` directly on mask operands — their implementations will be -> extended to accept mask types (treated as a per-lane bit-wise boolean op on the -> predicate). This also covers the `pnot`-style predicate complement needed by MERGE -> emulation (see [Appendix C](10-appendices.md)). +--- + +## Mask Boolean Ops (`vand` / `vor` / `vxor` / `vnot` on masks) + +The elementwise bitwise ops are reused directly on mask operands, treated as a +per-lane bit-wise boolean op on the predicate. + +- **example:** + ```mlir + // Predicate boolean ops on masks + %and = pto.vmi.vand %lt, %gt + : !pto.vmi.mask<128xpred>, !pto.vmi.mask<128xpred> + -> !pto.vmi.mask<128xpred> + %or = pto.vmi.vor %lt, %gt + : !pto.vmi.mask<128xpred>, !pto.vmi.mask<128xpred> + -> !pto.vmi.mask<128xpred> + %xor = pto.vmi.vxor %lt, %gt + : !pto.vmi.mask<128xpred>, !pto.vmi.mask<128xpred> + -> !pto.vmi.mask<128xpred> + %not = pto.vmi.vnot %lt + : !pto.vmi.mask<128xpred> -> !pto.vmi.mask<128xpred> + ``` diff --git a/docs/isa/vmi-isa/09-data-rearrange.md b/docs/isa/vmi-isa/09-data-rearrange.md index 5779151e2c..cde3190871 100644 --- a/docs/isa/vmi-isa/09-data-rearrange.md +++ b/docs/isa/vmi-isa/09-data-rearrange.md @@ -7,6 +7,7 @@ > has the same `L` and `T` as the inputs. Commonly used for real+imaginary and > value+index interleaving within a single vector register. + --- ## `pto.vmi.vintlv` @@ -58,6 +59,7 @@ -> !pto.vmi.vreg<64×f32>, !pto.vmi.vreg<64×f32> ``` + --- ## `pto.vmi.vdintlv` From 8d30be938df49c035493824dcfedd4bc12500d1f Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 16 Jul 2026 17:39:19 +0800 Subject: [PATCH 043/122] tci and trowexpand --- ...oas-implicit-tmp-materialization-design.md | 579 ++++++++++++++++++ include/PTO/Transforms/Passes.h | 2 + include/PTO/Transforms/Passes.td | 16 + lib/PTO/IR/PTO.cpp | 208 ++++++- lib/PTO/Transforms/CMakeLists.txt | 1 + .../Transforms/PTOMaterializeImplicitTmp.cpp | 290 +++++++++ test/lit/pto/easy_param_completion_emitc.pto | 4 +- test/lit/pto/tci_i16_emitc.pto | 4 +- .../pto/tci_implicit_tmp_level3_invalid.pto | 13 + .../pto/tci_implicit_tmp_materialization.pto | 25 + test/lit/pto/tci_tmp_contract_a3_invalid.pto | 27 + test/lit/pto/tci_ui16_emitc.pto | 4 +- test/lit/pto/tci_ui32_emitc.pto | 4 +- ...trowexpand_implicit_tmp_level3_invalid.pto | 24 + ...rowexpand_implicit_tmp_materialization.pto | 42 ++ .../trowexpand_tmp_contract_a3_invalid.pto | 33 + tools/ptoas/ptoas.cpp | 3 + 17 files changed, 1263 insertions(+), 16 deletions(-) create mode 100644 docs/designs/ptoas-implicit-tmp-materialization-design.md create mode 100644 lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp create mode 100644 test/lit/pto/tci_implicit_tmp_level3_invalid.pto create mode 100644 test/lit/pto/tci_implicit_tmp_materialization.pto create mode 100644 test/lit/pto/tci_tmp_contract_a3_invalid.pto create mode 100644 test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto create mode 100644 test/lit/pto/trowexpand_implicit_tmp_materialization.pto create mode 100644 test/lit/pto/trowexpand_tmp_contract_a3_invalid.pto diff --git a/docs/designs/ptoas-implicit-tmp-materialization-design.md b/docs/designs/ptoas-implicit-tmp-materialization-design.md new file mode 100644 index 0000000000..40ebab49f3 --- /dev/null +++ b/docs/designs/ptoas-implicit-tmp-materialization-design.md @@ -0,0 +1,579 @@ +# PTOAS Implicit Tmp Materialization Design + +## 背景 + +PTOAS 前端 IR 中很多 tile op 的 `tmp` operand 是可选的。当前如果用户没有显式写 `tmp`，PTOAS 会继续 lowering 到后端不带 `tmp` 的 C++ 接口。 + +后续希望改成：前端仍允许省略 `tmp`，但 PTOAS 在内部为需要 tmp-aware 后端接口的 op 自动补充合法 tmp tile，并让这些 tmp tile 和其它 tile buffer 一起进入 memplan 做 local addr 规划。 + +本文档的目标是给所有类似 `pto.tci` 的 optional tmp op 提供整体改造方案。每个 op 再根据自己的后端 tmp 规格，补充 op-specific 的 tmp requirement、MemoryEffects、verifier 和测试。 + +## 目标 + +- 保持前端 IR 兼容：用户仍然可以写不带 `tmp` 的目标 op。 +- 在 memplan 之前补齐隐式 tmp，使 tmp 作为普通 local allocation root 参与地址规划。 +- 对需要 tmp-aware 后端接口的 op，EmitC lowering 统一走带 tmp 的 C++ overload，避免继续选择 no-tmp 后端接口。 +- 每个 op 的 tmp shape、dtype、address space、layout、容量等约束由该 op 的后端接口规格决定。 +- 不在 EmitC lowering 中临时分配 tmp 地址。 +- 不在 memplan 中特殊创建 tmp；memplan 只负责规划已经存在的 root。 + +## 非目标 + +- 本阶段不一次性覆盖所有 optional tmp op。 +- 本阶段不改变用户显式提供 tmp 的语义。 +- 本阶段不为 level3 自动分配 tmp 地址。 +- 本阶段不引入新的全局 workspace 规划。 +- 本阶段不把所有 op 的 tmp 规格抽象成完全统一的 shape；不同 op 可以有不同 tmp requirement。 + +## 总体方案 + +新增一个 IR 规范化 pass： + +```text +pto-materialize-implicit-tmp +``` + +该 pass 运行在 `PTOViewToMemref` 之后、`pto-plan-memory` 之前： + +```text +PTOViewToMemref + -> pto-materialize-implicit-tmp + -> pto-plan-memory + -> PTOResolveReservedBuffers + -> sync passes + -> PTOMaterializeTileHandles + -> PTOToEmitC +``` + +pass 的职责是扫描所有已纳入改造的目标 op。如果 op 没有 tmp operand，就根据该 op 的 `TmpRequirement` 在 op 前插入 `pto.alloc_tile(no addr)`，并重写原 op，使其显式携带 tmp。 + +抽象流程： + +```text +target_op(no tmp) + -> lookup TmpRequirement(target_op) + -> create pto.alloc_tile(no addr) tmp + -> rewrite target_op(with tmp) + -> memplan assigns addr to tmp + -> EmitC sees tmp and emits tmp-aware overload +``` + +`TmpRequirement` 至少应包含： + +```text +AddressSpace space; +Type elementType; +StaticShape or MinBytes requirement; +Layout/layout-family requirement; +uint64_t minBytes; +bool requireExplicitAtLevel3; +``` + +对自动生成的 tmp： + +- 使用 tile-native `pto.alloc_tile(no addr)`。 +- 不创建 `memref.alloc`。 +- 不创建 `pto.pointer_cast` / `pto.bind_tile`。 +- 不设置 `addr`，由 memplan 统一规划。 +- 尽量使用静态 full-valid shape，即 `v_row/v_col` 与 `rows/cols` 一致，不额外携带 `valid_row` / `valid_col` operand。 +- 定义位置必须支配目标 op。 +- 生命周期由后续 liveness/memplan 根据真实 use 计算。 + +## Memplan 接入 + +自动生成的 tmp 是 tile-native `pto.alloc_tile(no addr)`，因此复用当前 memplan 路径： + +```text +pto.alloc_tile(no addr) + -> local allocation root + -> legacy/modern memplan 分配 offset + -> pto.alloc_tile addr = ... +``` + +legacy memplan 和 modern memplan 都应把自动生成的 tmp 当成普通 local allocation root。memplan 不应该知道“这是某个 op 的隐式 tmp”，也不应该在内部临时创建 tmp。 + +memplan 侧需要依赖 op 的 MemoryEffects / semantic no-alias 信息保证正确复用： + +- tmp 如果是 scratch buffer，应通过 `Write(tmp)` 建模，使 scratch-output conflict 能禁止 tmp 与同 op output 错误复用。 +- 如果某个 op 的 tmp 与 output 不能 alias，但 tmp 不适合建模成 scratch write，则应在 semantic no-alias side table 中显式加入 `forbidAlias(tmp, output)`。 +- 每个 op 的专项改造必须说明 tmp 和 output、input 之间的 alias 约束。 + +## Level 行为 + +### level1 / level2 + +level1/level2 下 memplan 会运行，因此允许省略 tmp： + +```text +target_op(no tmp) + -> pto-materialize-implicit-tmp + -> pto.alloc_tile(no addr) tmp + -> pto-plan-memory 补 addr +``` + +用户显式提供 tmp 时，仍需满足该 op 的 tmp verifier 约束。level1/level2 下用户不应显式指定 local addr，地址由 memplan 统一规划。 + +### level3 + +level3 下用户显式管理 local 地址，memplan 通常跳过。因此不应自动创建无地址 tmp。 + +通用规则： + +```text +level3 + target_op(no tmp) => pass/verifier 报错 +``` + +用户在 level3 使用已纳入改造的目标 op 时，必须显式提供合法 tmp，并保证 tmp 自身带合法 local addr，或满足现有 level3 显式地址规则。 + +诊断信息示例： + +```text + requires explicit tmp when compiling at level3 because PlanMemory is skipped +``` + +## EmitC Lowering + +目标 op 的 EmitC lowering 应保持简单： + +- `op.getTmp()` 为空：生成 no-tmp C++ 调用，或者在该 op 改造完成后仅作为未经过 materialize pass 的兜底路径。 +- `op.getTmp()` 非空：生成带 tmp 的 C++ 调用。 + +引入 `pto-materialize-implicit-tmp` 后，level1/level2 的目标 op 在 EmitC 前都会携带 tmp，因此会自然走带 tmp 的 overload。 + +不建议在 PTOToEmitC 中补 tmp，原因： + +- EmitC 阶段已经错过 memplan。 +- 临时生成 tmp 无法获得 local addr。 +- 会绕过 liveness、sync 和 semantic no-alias 分析。 + +## TCI 针对性改造 + +本节描述 `pto.tci` 作为第一批目标 op 的具体落地规则。后续其它 optional tmp op 应新增类似小节，分别说明自己的 tmp 规格、pass 行为、MemoryEffects、verifier 和测试计划。 + +### TCI Tmp 约束 + +`pto.tci` 当前 ODS 已经支持可选 tmp： + +```td +Optional:$tmp +``` + +PTOToEmitC 也已经根据 `op.getTmp()` 选择带 tmp 或不带 tmp 的 C++ 调用。因此 TCI 改造不需要改 `pto.tci` 的 IR 语法，关键是保证进入 EmitC 前缺省 tmp 已经被显式 materialize。 + +TCI 后端 C++ 接口存在两类 overload： + +```cpp +TCI(dst, start) +TCI(dst, start, tmp) +``` + +A2/A3 上 no-tmp overload 可能走 scalar loop；带 tmp overload 才能走更优路径。A5 接受 tmp，但 tmp 可以作为兼容占位，不额外引入有效计算约束。 + +TCI tmp 不应要求固定 shape，应按 PTO-ISA 文档中的精细化 tmp 约束校验容量。PTOAS 对用户显式 tmp 和自动生成 tmp 采用同一组 A2/A3 合法性规则： + +```text +loc = vec +dtype = 4-byte type: f32 / i32 / ui32 +shape = static shape +layout = row_major +fractal = 512 +capacity = product(shape) * sizeof(dtype) +``` + +A2/A3 的最小容量由 dst 元素类型决定： + +```text +b32 dst: i32 / ui32 -> tmp capacity >= 768 bytes +b16 dst: i16 / ui16 -> tmp capacity >= 1792 bytes +``` + +其中 `shape` 可以是任意静态形状，只要总容量满足对应 dst 类型的最小容量。例如 b32 dst 可以使用 `1x192xf32`，b16 dst 可以使用 `1x448xf32`。`Tile` 是 PTO-ISA 文档中推荐的方便形状无关分配，容量为 2048 bytes (2KiB)，可以同时覆盖 b32/b16。 + +A5 上 `tmp` Tile 被接受但不使用；A5 硬件直接使用 `vci` 向量指令，无需临时缓冲区。因此 A5 下 `pto.tci(no tmp)` 不需要自动 materialize tmp，用户显式传 tmp 时也不按 A2/A3 的容量规则校验。 + +### Pass 行为 + +对每个 `pto.tci`： + +- 如果已经有 `tmp`，pass 不修改。 +- A5 如果没有 `tmp`，pass 不修改；A5 后端直接使用 `vci`，不需要 tmp。 +- A2/A3 如果没有 `tmp`，且当前 build level 会运行 memplan，则自动补 tmp。 +- A2/A3 如果没有 `tmp`，但当前 level3 会跳过 memplan，则报错，要求用户显式提供带地址的 tmp。 + +重写前： + +```mlir +pto.tci ins(%s : i32) + outs(%dst : !pto.tile_buf) +``` + +A2/A3 重写后，以下以 b32 dst 自动生成 `f32 1x192` tmp 为例： + +```mlir +%tmp = pto.alloc_tile + : !pto.tile_buf + +pto.tci ins(%s, %tmp : i32, + !pto.tile_buf) + outs(%dst : !pto.tile_buf) +``` + +随后 memplan 会把 `%tmp` 当成普通 tile-native local allocation root，和其它 `pto.alloc_tile(no addr)` 一起规划 local address： + +```mlir +%tmp = pto.alloc_tile addr = %c4096_i64 + : !pto.tile_buf +``` + +TCI rewrite 需要保留原 op 的： + +- scalar operand `S`。 +- dst operand。 +- `descending` attr。 +- location。 +- 其它已有属性。 + +### MemoryEffects + +当前 `TCIOp::getEffects()` 只建模为： + +```text +Write(dst) +``` + +A2/A3 自动补 tmp 后，应改为： + +```text +Read(tmp) if tmp exists +Write(tmp) if tmp exists +Write(dst) +``` + +A5 上 tmp 被接受但不使用，因此不应把 tmp 建模为 Read/Write： + +```text +Write(dst) +``` + +原因： + +- liveness 需要看到 tmp 在 `pto.tci` 被使用。 +- sync pass 需要知道 `pto.tci` 会读 tmp 地址。 +- memplan 需要把 tmp 识别为 scratch buffer，避免 tmp 和同 op 的 dst 错误复用。 +- modern memplan 的 op semantic no-alias 和 root use 传播需要真实 use 信息。 + +如果 tmp 不被建模为 Read/Write，tmp 可能被认为没有 use 或不是 scratch，导致生命周期、复用或同步分析不准确。 + +TCI 也可以在 semantic no-alias side table 中显式加入： + +```text +op = pto.tci +forbidAlias(tmp, dst) +``` + +这不是 scratch conflict 生效的必要条件；A2/A3 只要 `TCIOp::getEffects()` 建模了 `Write(tmp)`，tmp 就会进入 scratch buffer conflict。但显式 side table 能防止未来有人调整 MemoryEffects 后破坏 tmp/dst no-alias 语义。 + +### Verifier + +`TCIOp::verify()` 应检查 tmp 是合法 tile buf，并满足后端接口容量约束： + +```text +如果 tmp 存在： + A5: tmp 被接受但不使用，不执行 A2/A3 tmp 容量校验。 + A2/A3: tmp 必须是 vec tile。 + A2/A3: tmp element type 必须是 4 字节类型（f32 / i32 / ui32）。 + A2/A3: tmp shape 必须是静态 shape。 + A2/A3: tmp layout 必须满足后端 TCI tmp 接口要求。 + A2/A3 b32 dst: tmp capacity 必须大于等于 768 bytes。 + A2/A3 b16 dst: tmp capacity 必须大于等于 1792 bytes。 +``` + +这里的关键是“容量满足接口约束”，而不是“shape 必须等于某个固定值”。TCI 按 dst 元素类型精细化检查： + +```text +// b32 dst: 768B 即可 +!pto.tile_buf // 合法 +!pto.tile_buf // 非法，容量不足 + +// b16 dst: 1792B 即可 +!pto.tile_buf // 合法 +!pto.tile_buf // 非法，容量不足 +``` + +自动生成 tmp 选择满足最小容量的 canonical shape： + +- b32 dst: `f32 1x192`。 +- b16 dst: `f32 1x448`。 +- A5: 不自动生成 tmp。 + +### 测试计划 + +#### lit：自动补 tmp + +新增用例： + +```text +test/lit/pto/tci_implicit_tmp_materialization.pto +``` + +检查： + +```text +CHECK: pto.alloc_tile +CHECK-SAME: dtype=f32 +CHECK: pto.tci ins(%{{.*}}, %{{.*}} +CHECK-NOT: memref.alloc +``` + +可以额外检查自动生成 shape：b32 dst 为 `f32 1x192`，b16 dst 为 `f32 1x448`。同时应覆盖 A5 下不自动生成 tmp。 + +#### lit：memplan 回写 addr + +检查 plan memory 后： + +```text +CHECK: pto.alloc_tile addr = +CHECK: pto.tci ins(%{{.*}}, %{{.*}} +``` + +legacy 和 modern 都应覆盖： + +```text +// RUN: ptoas --pto-level=level2 --plan-memory-impl=legacy ... +// RUN: ptoas --pto-level=level2 --plan-memory-impl=modern ... +``` + +#### lit：EmitC 走带 tmp overload + +检查 C++ 输出： + +```text +CHECK: TCI< +CHECK-SAME: Tile< +CHECK-SAME: float +CHECK: TCI{{.*}}({{.*}}, {{.*}}, {{.*}}) +``` + +#### lit：level3 负例 + +```text +level3 + pto.tci(no tmp) +``` + +期望： + +```text +expected-error {{pto.tci requires explicit tmp when compiling at level3}} +``` + +#### lit：verifier 负例 + +用户显式提供非法 tmp： + +- 非 vec space。 +- 非 f32 dtype。 +- dynamic shape。 +- layout 不满足 TCI tmp 接口约束。 +- A2/A3 b32 dst 的 tmp capacity 小于 768 bytes。 +- A2/A3 b16 dst 的 tmp capacity 小于 1792 bytes。 + +期望 verifier 报错。 + +## TROWEXPAND 二元 op 针对性改造 + +本节覆盖以下 row-expand 二元 op： + +```text +pto.trowexpandadd +pto.trowexpandsub +pto.trowexpandmul +pto.trowexpanddiv +pto.trowexpandmax +pto.trowexpandmin +``` + +这些 op 的 PTO-ISA 文档对 tmp 的描述一致：带 `TileDataTmp &tmp` 的 C++ overload 仅支持模式 1；A2/A3 上 tmp 用作行广播缓冲区；A5 接受 tmp 但不使用。 + +### RowExpand Tmp 约束 + +这些 op 有两种 row-broadcast 模式： + +- 模式 1：扩展操作数为 `ColMajor`，每行一个标量。带 tmp overload 仅支持该模式。 +- 模式 2：扩展操作数为 `RowMajor`，每行一个 32 字节块。该模式不需要 tmp，不应为了 tmp-aware overload 强行改写。 + +A2/A3 模式 1 下，tmp 作为 `vbrcb` 广播缓冲区使用。扩展操作数的每行标量会广播成一个 32 字节块；`vbrcb` repeat stride 为 8 个块，即 256 字节，每个 repeat 处理 8 行。 + +tmp 最小容量由 `R = dst.validRow` 决定： + +```text +if R < 256: + tmpBytes = ceil(R / 8) * 256 +else: + tmpBytes = 30 * 256 = 7680 +``` + +说明： + +- 当 `R >= 256` 时，后端按循环处理，每次循环最多 30 个 repeat，也就是 240 行；tmp 在循环间复用，因此每次循环只需要 7680 字节。 +- 一个紧凑的形状无关上界是 8KB，即 8192 字节。该上界可作为自动 materialize 的保守 canonical tmp 大小。 +- 不带 tmp 的 3 参数 overload 支持模式 1 和模式 2；对 A2/A3 的模式 1，后端使用内部 8KB 缓冲区 `TMP_UB_OFFSET`；模式 2 不需要广播缓冲区。 +- A5 硬件通过 `vlds` 广播模式原生支持行广播，tmp 被接口接受但不使用。 + +PTOAS 对用户显式 tmp 的合法性规则： + +```text +A2/A3: + op 必须是模式 1，才能使用显式 tmp。 + tmp 必须是 vec tile。 + tmp shape 必须静态可计算容量，或后续 verifier 能证明容量满足公式。 + tmp capacity >= min(ceil(R / 8) * 256, 7680)。 + +A5: + tmp 被接受但不使用，不执行 A2/A3 tmp 容量校验。 +``` + +### Pass 行为 + +对每个目标 row-expand 二元 op： + +- 如果已经有 `tmp`，pass 不修改，但 verifier 需要保证它只用于合法模式。 +- A5 如果没有 `tmp`，pass 不修改。 +- A2/A3 如果没有 `tmp`，且 op 是模式 1、当前 build level 会运行 memplan，则自动补 tmp。 +- A2/A3 如果没有 `tmp`，op 是模式 1、但当前 level3 会跳过 memplan，则报错，要求用户显式提供带地址的 tmp。 +- 模式 2 不需要 tmp；pass 不应自动补 tmp，也不应强制改成带 tmp overload。 + +自动补 tmp 的 canonical shape 建议采用形状无关上界： + +```mlir +%tmp = pto.alloc_tile + : !pto.tile_buf, + rows=1, cols=<8192 / sizeof(dst element type)>, + v_row=1, v_col=<8192 / sizeof(dst element type)>, + blayout=row_major, slayout=none_box, + fractal=512, pad=0> +``` + +默认使用 `dst element type` 作为 tmp element type，以贴合 row-expand 后端模板参数；如果后续确认某些后端实现允许更宽松的 tmp dtype，可在对应 op-specific verifier 中放宽。 + +这样不需要在 materialize pass 中依赖 `dst.validRow` 是否为静态值，也能覆盖 A2/A3 模式 1 的最大每轮 tmp 需求。后续如果希望节省 UB，可以在能静态证明 `R` 时生成更小 tmp： + +```text +tmpBytes = min(ceil(R / 8) * 256, 7680) +``` + +### MemoryEffects + +A2/A3 上这些 op 的 tmp 是广播 scratch buffer，应该建模为： + +```text +Read(non-tmp inputs) +Read(tmp) if tmp exists +Write(tmp) if tmp exists +Write(dst) +``` + +其中 `Write(tmp)` 用于让 memplan 的 scratch-output conflict 禁止 tmp 与同 op 的 `dst` 错误复用。 + +A5 上 tmp 被接受但不使用，因此不应把 tmp 建模为 Read/Write： + +```text +Read(non-tmp inputs) +Write(dst) +``` + +如果未来某个 row-expand op 的 MemoryEffects 不适合用 `Write(tmp)` 表达 scratch 语义，也应在 semantic no-alias side table 中显式加入： + +```text +op = pto.trowexpand* +forbidAlias(tmp, dst) +``` + +### Verifier + +这些 op 的 verifier 需要区分模式和 arch： + +```text +如果 tmp 存在： + A5: tmp 被接受但不使用，不执行 A2/A3 tmp 容量校验。 + A2/A3: op 必须是模式 1，即扩展操作数为 ColMajor 每行标量。 + A2/A3: tmp 必须是 vec tile。 + A2/A3: tmp capacity 必须满足 min(ceil(dst.validRow / 8) * 256, 7680)。 +``` + +模式识别规则沿用 ISA 文档： + +- `src0` 或 `src1` 中恰好一个与 `dst` 有相同 valid shape，该 operand 是全尺寸操作数。 +- 另一个 operand 是扩展操作数。 +- 扩展操作数为 `ColMajor` 且每行一个标量时是模式 1。 +- 扩展操作数为 `RowMajor` 且每行 `32 / sizeof(T)` 列时是模式 2。 + +如果 `dst.validRow` 是动态值，verifier 无法精确证明用户显式 tmp 是否足够小时，可以采用保守规则： + +- 用户显式 tmp 至少 8192 字节；或 +- 后续引入运行时/符号约束证明 tmp capacity 满足公式。 + +自动生成 tmp 建议先使用 8192 字节 canonical 上界，因此不会受动态 `dst.validRow` 影响。 + +### 测试计划 + +#### lit：自动补 tmp + +为至少一个代表 op 增加 A2/A3 模式 1 用例，例如 `pto.trowexpandadd(no tmp)`： + +```text +CHECK: pto.alloc_tile +CHECK: pto.trowexpandadd ins(%{{.*}}, %{{.*}}, %{{.*}} +CHECK-NOT: memref.alloc +``` + +同时检查 A5 下不自动生成 tmp。 + +#### lit：模式 2 不补 tmp + +构造 RowMajor 扩展操作数的模式 2 用例，确认 pass 不自动补 tmp，并继续走 no-tmp overload。 + +#### lit：memplan 回写 addr + +检查自动生成 tmp 在 plan memory 后带 `addr`： + +```text +CHECK: pto.alloc_tile addr = +CHECK: pto.trowexpand{{.*}} ins(%{{.*}}, %{{.*}}, %{{.*}} +``` + +legacy 和 modern 都应覆盖。 + +#### lit：level3 负例 + +A2/A3 level3 + 模式 1 + no tmp 应报错： + +```text +expected-error {{requires explicit tmp when compiling at level3}} +``` + +A5 level3 + no tmp 不应因为 tmp 缺失报错。 + +#### lit：verifier 负例 + +需要覆盖： + +- A2/A3 显式 tmp 用在模式 2，报错。 +- A2/A3 显式 tmp capacity 小于公式要求，报错。 +- A2/A3 动态 `dst.validRow` 且显式 tmp 小于 8192 字节，按保守规则报错。 +- A5 显式 tmp 不触发 A2/A3 容量校验。 + +## 后续扩展 + +后续新增其它 optional tmp op 时，需要补充一个 op-specific 小节，并明确： + +- op 的 tmp 后端接口规格。 +- 自动生成 tmp 的 canonical shape。 +- 用户显式 tmp 的 verifier 规则。 +- tmp 的 MemoryEffects。 +- tmp 与 input/output 的 alias 约束。 +- level3 下是否要求显式 tmp。 +- lit 覆盖自动补 tmp、memplan 回写 addr、EmitC overload、level3 负例和 verifier 负例。 diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index c1d9e82a14..a9b82c8e3f 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -75,6 +75,8 @@ createPlanMemoryModernPass(const PlanMemoryOptions &options); std::unique_ptr createPTORemoveRedundantBarrierPass(); std::unique_ptr createPTOValidateIntToPtrUsesPass(); std::unique_ptr createPTORematerializeFixpipeVectorQuantPass(); +std::unique_ptr +createPTOMaterializeImplicitTmpPass(bool requireExplicitTmp = false); std::unique_ptr createPTOResolveBufferSelectPass(); std::unique_ptr createInferPTOLayoutPass(); std::unique_ptr createPTOA5NormalizeTMovPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 7b08319dda..97206bb0fb 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -192,6 +192,22 @@ def PTORematerializeFixpipeVectorQuant let dependentDialects = ["mlir::pto::PTODialect", "mlir::func::FuncDialect"]; } +def PTOMaterializeImplicitTmp + : Pass<"pto-materialize-implicit-tmp", "func::FuncOp"> { + let summary = "Materialize implicit tmp tiles for PTO ops before memplan"; + let description = [{ + Rewrites PTO ops with optional tmp operands into explicit tmp forms when + the backend tmp-aware overload is required. The synthesized tmp is emitted + as tile-native `pto.alloc_tile(no addr)` so PlanMemory can assign its local + address together with other tile buffers. + }]; + let constructor = "mlir::pto::createPTOMaterializeImplicitTmpPass()"; + let dependentDialects = [ + "mlir::pto::PTODialect", + "mlir::func::FuncDialect" + ]; +} + def PlanMemory : Pass<"pto-plan-memory", "ModuleOp"> { let summary = "Plan memory for PTO Ops"; let constructor = "mlir::pto::createPlanMemoryPass()"; diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 86907e3255..90362862ab 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -2015,6 +2015,12 @@ static bool isRowMajorTileBuf(Type ty) { return tb && tb.getBLayoutValueI32() == static_cast(pto::BLayout::RowMajor); } +static bool isColMajorTileBuf(Type ty) { + auto tb = mlir::dyn_cast(ty); + return tb && tb.getBLayoutValueI32() == + static_cast(pto::BLayout::ColMajor); +} + static LogicalResult verifyRowReductionSrcLayout(Operation *op, Type ty, StringRef name) { if (failed(verifyTileBufCommon(op, ty, name))) @@ -5841,6 +5847,33 @@ LogicalResult pto::TCIOp::verify() { if (bw != 16 && bw != 32) return emitOpError("expects dst element type to be i16/i32"); + if (getTmp() && getTargetArch(getOperation()) != PTOArch::A5) { + auto tmpTy = mlir::dyn_cast(getTmp().getType()); + if (!tmpTy) + return emitOpError("expects tmp to be a tile buffer"); + auto tmpSpace = + mlir::dyn_cast_or_null(tmpTy.getMemorySpace()); + if (!tmpSpace || tmpSpace.getAddressSpace() != AddressSpace::VEC) + return emitOpError("expects tmp to be in vec address space"); + Type tmpElemTy = tmpTy.getElementType(); + if (!(tmpElemTy.isF32() || tmpElemTy.isInteger(32))) + return emitOpError("expects A2/A3 tmp element type to be a 4-byte type"); + if (tmpTy.getBLayoutValueI32() != static_cast(BLayout::RowMajor)) + return emitOpError("expects tmp blayout to be row_major"); + if (tmpTy.getSLayoutValueI32() != static_cast(SLayout::NoneBox)) + return emitOpError("expects tmp slayout to be none_box"); + if (tmpTy.getSFractalSizeI32() != 512) + return emitOpError("expects tmp fractal size to be 512"); + auto tmpBytes = getStaticByteSize(tmpTy); + if (!tmpBytes) + return emitOpError("expects tmp to have static byte size"); + uint64_t minTmpBytes = bw == 32 ? 768 : 1792; + if (*tmpBytes < minTmpBytes) + return emitOpError("expects A2/A3 tmp capacity to be at least ") + << minTmpBytes << " bytes for " << bw + << "-bit dst element type"; + } + auto sTy = mlir::dyn_cast(getOperand(0).getType()); if (!sTy) return emitOpError("expects S to be integer"); @@ -11723,6 +11756,112 @@ static FailureOr verifyTRowExpandBinaryCore(Operation *op, Type src0Ty, return getElemTy(src0Ty); } +enum class TRowExpandBinaryMode { + Unknown, + Mode1ColMajorScalar, + Mode2RowMajorBlock, +}; + +static bool validShapesCompatibleForTRowExpand(ArrayRef lhs, + ArrayRef rhs) { + if (lhs.size() != rhs.size()) + return false; + for (auto [l, r] : llvm::zip(lhs, rhs)) { + if (l != ShapedType::kDynamic && r != ShapedType::kDynamic && l != r) + return false; + } + return true; +} + +static TRowExpandBinaryMode classifyTRowExpandBinaryMode(Type src0Ty, + Type src1Ty, + Type dstTy) { + auto src0Valid = getValidShapeVec(src0Ty); + auto src1Valid = getValidShapeVec(src1Ty); + auto dstValid = getValidShapeVec(dstTy); + if (src0Valid.size() != 2 || src1Valid.size() != 2 || dstValid.size() != 2) + return TRowExpandBinaryMode::Unknown; + + Type expandedTy; + ArrayRef expandedValid; + if (validShapesCompatibleForTRowExpand(src0Valid, dstValid)) { + expandedTy = src1Ty; + expandedValid = src1Valid; + } else if (validShapesCompatibleForTRowExpand(src1Valid, dstValid)) { + expandedTy = src0Ty; + expandedValid = src0Valid; + } else { + return TRowExpandBinaryMode::Unknown; + } + + int64_t expandedCols = expandedValid[1]; + if (isColMajorTileBuf(expandedTy) && + (expandedCols == ShapedType::kDynamic || expandedCols == 1)) + return TRowExpandBinaryMode::Mode1ColMajorScalar; + + std::optional elemBytes = getElemBytes(getElemTy(dstTy)); + if (!elemBytes || *elemBytes == 0) + return TRowExpandBinaryMode::Unknown; + int64_t expectedMode2Cols = 32 / *elemBytes; + if (isRowMajorTileBuf(expandedTy) && + (expandedCols == ShapedType::kDynamic || + expandedCols == expectedMode2Cols)) + return TRowExpandBinaryMode::Mode2RowMajorBlock; + + return TRowExpandBinaryMode::Unknown; +} + +static int64_t getTRowExpandTmpMinBytes(int64_t dstValidRows) { + if (dstValidRows == ShapedType::kDynamic) + return 8192; + if (dstValidRows < 0) + return 8192; + if (dstValidRows < 256) + return ceilDivInt64(dstValidRows, 8) * 256; + return 30 * 256; +} + +static std::optional getStaticTileCapacityBytes(Type ty) { + auto numElems = getStaticNumElements(getShapeVec(ty)); + auto elemBytes = getElemBytes(getElemTy(ty)); + if (!numElems || !elemBytes) + return std::nullopt; + return *numElems * *elemBytes; +} + +static LogicalResult verifyTRowExpandImplicitTmpContract( + Operation *op, Type src0Ty, Type src1Ty, Type dstTy, Type tmpTy, + bool hasTmp, PTOArch targetArch) { + if (!hasTmp || targetArch == PTOArch::A5) + return success(); + + if (classifyTRowExpandBinaryMode(src0Ty, src1Ty, dstTy) != + TRowExpandBinaryMode::Mode1ColMajorScalar) { + return op->emitOpError( + "expects A2/A3 tmp-form trowexpand to use mode 1 " + "(ColMajor per-row scalar expanded operand)"); + } + + if (failed(verifyVecTileStorage(op, tmpTy, "tmp"))) + return failure(); + if (getElemTy(tmpTy) != getElemTy(dstTy)) + return op->emitOpError("expects tmp and dst to have the same element type"); + + auto dstValid = getValidShapeVec(dstTy); + if (dstValid.size() != 2) + return op->emitOpError("expects dst to have rank-2 valid_shape"); + int64_t minBytes = getTRowExpandTmpMinBytes(dstValid[0]); + std::optional tmpBytes = getStaticTileCapacityBytes(tmpTy); + if (!tmpBytes) + return op->emitOpError( + "expects A2/A3 trowexpand tmp capacity to be statically known"); + if (*tmpBytes < minBytes) + return op->emitOpError() + << "expects A2/A3 trowexpand tmp capacity to be at least " + << minBytes << " bytes, but got " << *tmpBytes << " bytes"; + return success(); +} + mlir::LogicalResult mlir::pto::TRowExpandDivOp::verify() { auto verifyByArch = [&](PTOArch targetArch) -> LogicalResult { Type src0Ty = getSrc0().getType(); @@ -11746,6 +11885,11 @@ mlir::LogicalResult mlir::pto::TRowExpandDivOp::verify() { } if (getPrecisionType() == pto::DivPrecision::HighPrecision && !getTmp()) return emitOpError("expects tmp when precisionType is high_precision"); + if (failed(verifyTRowExpandImplicitTmpContract( + getOperation(), src0Ty, src1Ty, dstTy, + getTmp() ? getTmp().getType() : Type{}, static_cast(getTmp()), + targetArch))) + return failure(); return mlir::success(); }; auto verifyA2A3 = [&]() -> LogicalResult { return verifyByArch(PTOArch::A3); }; @@ -11775,6 +11919,11 @@ mlir::LogicalResult mlir::pto::TRowExpandMulOp::verify() { return emitOpError( "expects A2/A3 trowexpandmul element type to be i16/i32/f16/f32"); } + if (failed(verifyTRowExpandImplicitTmpContract( + getOperation(), src0Ty, src1Ty, dstTy, + getTmp() ? getTmp().getType() : Type{}, static_cast(getTmp()), + targetArch))) + return failure(); return mlir::success(); }; auto verifyA2A3 = [&]() -> LogicalResult { return verifyByArch(PTOArch::A3); }; @@ -11804,6 +11953,11 @@ mlir::LogicalResult mlir::pto::TRowExpandSubOp::verify() { return emitOpError( "expects A2/A3 trowexpandsub element type to be i16/i32/f16/f32"); } + if (failed(verifyTRowExpandImplicitTmpContract( + getOperation(), src0Ty, src1Ty, dstTy, + getTmp() ? getTmp().getType() : Type{}, static_cast(getTmp()), + targetArch))) + return failure(); return mlir::success(); }; auto verifyA2A3 = [&]() -> LogicalResult { return verifyByArch(PTOArch::A3); }; @@ -11855,6 +12009,11 @@ mlir::LogicalResult mlir::pto::TRowExpandAddOp::verify() { if (src1Col != ShapedType::kDynamic && src1Col != 1) return emitOpError("expects non-row-major src1 valid_shape[1] to be 1"); } + if (failed(verifyTRowExpandImplicitTmpContract( + getOperation(), src0Ty, src1Ty, dstTy, + getTmp() ? getTmp().getType() : Type{}, static_cast(getTmp()), + targetArch))) + return failure(); return mlir::success(); }; auto verifyA2A3 = [&]() -> LogicalResult { return verifyByArch(PTOArch::A3); }; @@ -11866,6 +12025,7 @@ static LogicalResult verifyTRowExpandReduceLikeOp(Operation *op, Type src0Ty, Type src1Ty, Type dstTy, Type tmpTy, bool hasTmp, PTOArch targetArch, + bool enforceTmpContract, StringRef opName, bool allowIntegerTypes) { if (failed(verifyTileBufCommon(op, src0Ty, "src0")) || @@ -11984,14 +12144,23 @@ static LogicalResult verifyTRowExpandReduceLikeOp(Operation *op, Type src0Ty, // (A5 tmp-form invariant is checked earlier, before the empty-marker accept.) + auto verifyTmpContract = [&]() -> LogicalResult { + if (!enforceTmpContract) + return success(); + return verifyTRowExpandImplicitTmpContract(op, src0Ty, src1Ty, dstTy, + tmpTy, hasTmp, targetArch); + }; + if (src0MatchesDst) { if (succeeded(checkFullAndBroadcast(src0Ty, src0Valid, "src0", src1Ty, - src1Valid, "src1"))) + src1Valid, "src1")) && + succeeded(verifyTmpContract())) return success(); } if (src1MatchesDst) { if (succeeded(checkFullAndBroadcast(src1Ty, src1Valid, "src1", src0Ty, - src0Valid, "src0"))) + src0Valid, "src0")) && + succeeded(verifyTmpContract())) return success(); } @@ -12005,6 +12174,7 @@ mlir::LogicalResult mlir::pto::TRowExpandExpdifOp::verify() { getSrc1().getType(), getDst().getType(), getTmp() ? getTmp().getType() : Type{}, (bool)getTmp(), PTOArch::A3, + /*enforceTmpContract=*/false, "trowexpandexpdif", /*allowIntegerTypes=*/false); }; @@ -12013,6 +12183,7 @@ mlir::LogicalResult mlir::pto::TRowExpandExpdifOp::verify() { getSrc1().getType(), getDst().getType(), getTmp() ? getTmp().getType() : Type{}, (bool)getTmp(), PTOArch::A5, + /*enforceTmpContract=*/false, "trowexpandexpdif", /*allowIntegerTypes=*/false); }; @@ -12025,6 +12196,7 @@ mlir::LogicalResult mlir::pto::TRowExpandMaxOp::verify() { getSrc1().getType(), getDst().getType(), getTmp() ? getTmp().getType() : Type{}, (bool)getTmp(), PTOArch::A3, + /*enforceTmpContract=*/true, "trowexpandmax", /*allowIntegerTypes=*/true); }; @@ -12033,6 +12205,7 @@ mlir::LogicalResult mlir::pto::TRowExpandMaxOp::verify() { getSrc1().getType(), getDst().getType(), getTmp() ? getTmp().getType() : Type{}, (bool)getTmp(), PTOArch::A5, + /*enforceTmpContract=*/true, "trowexpandmax", /*allowIntegerTypes=*/true); }; @@ -12045,6 +12218,7 @@ mlir::LogicalResult mlir::pto::TRowExpandMinOp::verify() { getSrc1().getType(), getDst().getType(), getTmp() ? getTmp().getType() : Type{}, (bool)getTmp(), PTOArch::A3, + /*enforceTmpContract=*/true, "trowexpandmin", /*allowIntegerTypes=*/true); }; @@ -12053,6 +12227,7 @@ mlir::LogicalResult mlir::pto::TRowExpandMinOp::verify() { getSrc1().getType(), getDst().getType(), getTmp() ? getTmp().getType() : Type{}, (bool)getTmp(), PTOArch::A5, + /*enforceTmpContract=*/true, "trowexpandmin", /*allowIntegerTypes=*/true); }; @@ -14181,6 +14356,11 @@ PTO_DEFINE_UNARY_EFFECTS(TAndSOp, getSrcMutable(), getDstMutable()) // TCI: Write(dst) (generates sequence) void TCIOp::getEffects( SmallVectorImpl> &effects) { + if (auto tmp = getTmpMutable(); + !tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14449,8 +14629,10 @@ void TRowExpandDivOp::getEffects( PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); auto tmp = getTmpMutable(); - if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14459,8 +14641,10 @@ void TRowExpandMulOp::getEffects( PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); auto tmp = getTmpMutable(); - if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14469,8 +14653,10 @@ void TRowExpandSubOp::getEffects( PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); auto tmp = getTmpMutable(); - if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14479,8 +14665,10 @@ void TRowExpandAddOp::getEffects( PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); auto tmp = getTmpMutable(); - if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14499,8 +14687,10 @@ void TRowExpandMaxOp::getEffects( PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); auto tmp = getTmpMutable(); - if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14509,8 +14699,10 @@ void TRowExpandMinOp::getEffects( PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); auto tmp = getTmpMutable(); - if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index 75f2c25f85..5dbabf5e90 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -75,6 +75,7 @@ add_mlir_dialect_library(PTOTransforms InsertSync/PTOInsertSync.cpp PTOInjectBarrierAllSync.cpp InsertSync/InsertSyncDebug.cpp + PTOMaterializeImplicitTmp.cpp PTORematerializeFixpipeVectorQuant.cpp PTOValidateIntToPtrUses.cpp InsertTemplateAttributes.cpp diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp new file mode 100644 index 0000000000..abb14a89e2 --- /dev/null +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -0,0 +1,290 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- PTOMaterializeImplicitTmp.cpp --------------------------------------===// + +#include "PTO/Transforms/Passes.h" + +#include "PTO/IR/PTO.h" +#include "PTO/IR/PTODialect.h" +#include "PTO/IR/PTOTypeUtils.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/TypeSwitch.h" + +using namespace mlir; + +namespace { + +static pto::TileBufConfigAttr makeRowMajorNoneBoxConfig(MLIRContext *ctx) { + OpBuilder builder(ctx); + return pto::TileBufConfigAttr::get( + ctx, pto::BLayoutAttr::get(ctx, pto::BLayout::RowMajor), + pto::SLayoutAttr::get(ctx, pto::SLayout::NoneBox), + builder.getI32IntegerAttr(512), + pto::PadValueAttr::get(ctx, pto::PadValue::Null), + pto::CompactModeAttr::get(ctx, pto::CompactMode::Null)); +} + +static unsigned getTCIDstBitWidth(pto::TCIOp op) { + auto tileTy = dyn_cast(op.getDst().getType()); + if (!tileTy) + return 0; + auto elemTy = dyn_cast(tileTy.getElementType()); + if (!elemTy) + return 0; + return elemTy.getWidth(); +} + +static pto::TileBufType makeTCITmpType(MLIRContext *ctx, unsigned dstBitWidth) { + // PTO-ISA TCI A2/A3 vector path needs 768B for b32 dst and 1792B for + // b16 dst. Use an f32 1xN tmp with the exact minimum capacity. + int64_t cols = dstBitWidth == 16 ? 448 : 192; + return pto::TileBufType::get( + ctx, {1, cols}, Float32Type::get(ctx), + pto::AddressSpaceAttr::get(ctx, pto::AddressSpace::VEC), {1, cols}, + makeRowMajorNoneBoxConfig(ctx)); +} + +static std::optional getElemBytes(Type elemTy) { + unsigned bits = pto::getPTOStorageElemBitWidth(elemTy); + if (bits == 0 || bits % 8 != 0) + return std::nullopt; + return bits / 8; +} + +static SmallVector getValidShapeVec(Type ty) { + if (auto tileTy = dyn_cast(ty)) + return SmallVector(tileTy.getValidShape().begin(), + tileTy.getValidShape().end()); + return {}; +} + +static bool validShapesCompatible(ArrayRef lhs, + ArrayRef rhs) { + if (lhs.size() != rhs.size()) + return false; + for (auto [l, r] : llvm::zip(lhs, rhs)) { + if (l != ShapedType::kDynamic && r != ShapedType::kDynamic && l != r) + return false; + } + return true; +} + +static bool isRowMajorTile(Value value) { + auto tileTy = dyn_cast(value.getType()); + return tileTy && tileTy.getBLayoutValueI32() == + static_cast(pto::BLayout::RowMajor); +} + +static bool isColMajorTile(Value value) { + auto tileTy = dyn_cast(value.getType()); + return tileTy && tileTy.getBLayoutValueI32() == + static_cast(pto::BLayout::ColMajor); +} + +enum class RowExpandMode { + Unknown, + Mode1ColMajorScalar, + Mode2RowMajorBlock, +}; + +static RowExpandMode classifyTRowExpandBinaryMode(Value src0, Value src1, + Value dst) { + auto dstValid = getValidShapeVec(dst.getType()); + auto src0Valid = getValidShapeVec(src0.getType()); + auto src1Valid = getValidShapeVec(src1.getType()); + if (dstValid.size() != 2 || src0Valid.size() != 2 || src1Valid.size() != 2) + return RowExpandMode::Unknown; + + Value expanded; + ArrayRef expandedValid; + if (validShapesCompatible(src0Valid, dstValid)) { + expanded = src1; + expandedValid = src1Valid; + } else if (validShapesCompatible(src1Valid, dstValid)) { + expanded = src0; + expandedValid = src0Valid; + } else { + return RowExpandMode::Unknown; + } + + int64_t expandedCols = expandedValid[1]; + if (isColMajorTile(expanded) && + (expandedCols == ShapedType::kDynamic || expandedCols == 1)) + return RowExpandMode::Mode1ColMajorScalar; + + auto dstTileTy = dyn_cast(dst.getType()); + if (!dstTileTy) + return RowExpandMode::Unknown; + auto elemBytes = getElemBytes(dstTileTy.getElementType()); + if (!elemBytes || *elemBytes == 0) + return RowExpandMode::Unknown; + int64_t expectedMode2Cols = 32 / *elemBytes; + if (isRowMajorTile(expanded) && + (expandedCols == ShapedType::kDynamic || + expandedCols == expectedMode2Cols)) + return RowExpandMode::Mode2RowMajorBlock; + + return RowExpandMode::Unknown; +} + +static pto::TileBufType makeTRowExpandTmpType(MLIRContext *ctx, + pto::TileBufType dstTy) { + constexpr int64_t kTmpBytes = 8192; + std::optional elemBytes = getElemBytes(dstTy.getElementType()); + int64_t cols = elemBytes && *elemBytes > 0 ? kTmpBytes / *elemBytes : 2048; + return pto::TileBufType::get( + ctx, {1, cols}, dstTy.getElementType(), + pto::AddressSpaceAttr::get(ctx, pto::AddressSpace::VEC), {1, cols}, + makeRowMajorNoneBoxConfig(ctx)); +} + +static void replaceTRowExpandBinaryOpWithTmp(Operation *op, Value src0, + Value src1, Value tmp, Value dst) { + OpBuilder builder(op); + OperationState state(op->getLoc(), op->getName()); + state.addOperands({src0, src1, tmp, dst}); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr({1, 1, 1, 1})); + for (NamedAttribute attr : op->getAttrs()) { + if (attr.getName() == "operandSegmentSizes") + continue; + state.addAttribute(attr.getName(), attr.getValue()); + } + builder.create(state); + op->erase(); +} + +template +static LogicalResult materializeTRowExpandTmp(OpTy op, bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp()) + return success(); + if (pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5) + return success(); + + RowExpandMode mode = + classifyTRowExpandBinaryMode(op.getSrc0(), op.getSrc1(), op.getDst()); + if (mode != RowExpandMode::Mode1ColMajorScalar) + return success(); + + if (requireExplicitTmp) { + return op.emitOpError( + "requires explicit tmp for A2/A3 row-expand mode 1 when PlanMemory is skipped"); + } + + auto dstTy = dyn_cast(op.getDst().getType()); + if (!dstTy) + return op.emitOpError("expects tile_buf dst when materializing implicit tmp"); + + OpBuilder builder(op); + Value tmp = + builder + .create(op.getLoc(), + makeTRowExpandTmpType(ctx, dstTy), Value(), + Value(), Value()) + .getResult(); + replaceTRowExpandBinaryOpWithTmp(op.getOperation(), op.getSrc0(), op.getSrc1(), + tmp, op.getDst()); + return success(); +} + +struct PTOMaterializeImplicitTmpPass + : public PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PTOMaterializeImplicitTmpPass) + + PTOMaterializeImplicitTmpPass() = default; + explicit PTOMaterializeImplicitTmpPass(bool requireExplicitTmp) + : requireExplicitTmp(requireExplicitTmp) {} + + StringRef getArgument() const final { return "pto-materialize-implicit-tmp"; } + StringRef getDescription() const final { + return "Materialize implicit tmp tiles for PTO ops before memplan"; + } + + void runOnOperation() override { + func::FuncOp func = getOperation(); + MLIRContext *ctx = func.getContext(); + bool failed = false; + + SmallVector tciOps; + func.walk([&](pto::TCIOp op) { + if (!op.getTmp()) + tciOps.push_back(op); + }); + + for (pto::TCIOp op : tciOps) { + if (pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5) + continue; + + if (requireExplicitTmp) { + op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + failed = true; + continue; + } + + OpBuilder builder(op); + Location loc = op.getLoc(); + auto tmpType = makeTCITmpType(ctx, getTCIDstBitWidth(op)); + Value tmp = + builder.create(loc, tmpType, Value(), Value(), + Value()) + .getResult(); + + auto newOp = builder.create( + loc, TypeRange{}, op.getS(), tmp, op.getDst(), + op.getDescendingAttr()); + for (NamedAttribute attr : op->getAttrs()) { + if (attr.getName() == "operandSegmentSizes") + continue; + newOp->setAttr(attr.getName(), attr.getValue()); + } + op.erase(); + } + + SmallVector rowExpandOps; + func.walk([&](Operation *op) { + if (isa(op)) + rowExpandOps.push_back(op); + }); + + for (Operation *op : rowExpandOps) { + LogicalResult result = + llvm::TypeSwitch(op) + .Case( + [&](auto typedOp) { + return materializeTRowExpandTmp(typedOp, requireExplicitTmp, + ctx); + }) + .Default([](Operation *) { return success(); }); + if (mlir::failed(result)) + failed = true; + } + + if (failed) + signalPassFailure(); + } + +private: + bool requireExplicitTmp = false; +}; + +} // namespace + +std::unique_ptr +mlir::pto::createPTOMaterializeImplicitTmpPass(bool requireExplicitTmp) { + return std::make_unique(requireExplicitTmp); +} diff --git a/test/lit/pto/easy_param_completion_emitc.pto b/test/lit/pto/easy_param_completion_emitc.pto index 53636240bb..36c9f30b80 100644 --- a/test/lit/pto/easy_param_completion_emitc.pto +++ b/test/lit/pto/easy_param_completion_emitc.pto @@ -4,8 +4,8 @@ module { func.func @tci_with_tmp(%dst: !pto.partition_tensor_view<1x16xi16>) { %c0_i16 = arith.constant 0 : i16 %dst_tile = pto.alloc_tile : !pto.tile_buf - %tmp_tile = pto.alloc_tile : !pto.tile_buf - pto.tci ins(%c0_i16, %tmp_tile : i16, !pto.tile_buf) + %tmp_tile = pto.alloc_tile : !pto.tile_buf + pto.tci ins(%c0_i16, %tmp_tile : i16, !pto.tile_buf) outs(%dst_tile : !pto.tile_buf) pto.tstore ins(%dst_tile : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x16xi16>) {layout = #pto.layout, pto.inferred_layout = true} diff --git a/test/lit/pto/tci_i16_emitc.pto b/test/lit/pto/tci_i16_emitc.pto index cb7f1558c6..4fd5917dd8 100644 --- a/test/lit/pto/tci_i16_emitc.pto +++ b/test/lit/pto/tci_i16_emitc.pto @@ -10,5 +10,5 @@ module { } } -// A3: TCI, int16_t, 0>( -// A3-NOT: TCI, int32_t, 0>( +// A3: TCI, {{.*}}, int16_t, 0>({{.*}}, {{.*}}, {{.*}}) +// A3-NOT: TCI, {{.*}}, int32_t, 0>( diff --git a/test/lit/pto/tci_implicit_tmp_level3_invalid.pto b/test/lit/pto/tci_implicit_tmp_level3_invalid.pto new file mode 100644 index 0000000000..bdf628d5c1 --- /dev/null +++ b/test/lit/pto/tci_implicit_tmp_level3_invalid.pto @@ -0,0 +1,13 @@ +// RUN: not ptoas --pto-arch=a3 --pto-level=level3 %s 2>&1 | FileCheck %s + +module { + func.func @tci_implicit_tmp_level3() { + %c0_i32 = arith.constant 0 : i32 + %addr = arith.constant 0 : i64 + %tile = pto.alloc_tile addr = %addr : !pto.tile_buf + // CHECK: error: 'pto.tci' op requires explicit tmp when PlanMemory is skipped + pto.tci ins(%c0_i32 : i32) + outs(%tile : !pto.tile_buf) + return + } +} diff --git a/test/lit/pto/tci_implicit_tmp_materialization.pto b/test/lit/pto/tci_implicit_tmp_materialization.pto new file mode 100644 index 0000000000..78cec449e5 --- /dev/null +++ b/test/lit/pto/tci_implicit_tmp_materialization.pto @@ -0,0 +1,25 @@ +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a3 --pto-level=level2 %s 2>&1 | FileCheck %s --check-prefix=CPP +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5IR + +module { + func.func @tci_implicit_tmp(%dst: memref<32xi32, #pto.address_space>) { + %c0_i32 = arith.constant 0 : i32 + %src = memref.reinterpret_cast %dst to offset: [0], sizes: [1, 32], strides: [32, 1] {layout = #pto.layout} : memref<32xi32, #pto.address_space> to memref<1x32xi32, strided<[32, 1], offset: ?>, #pto.address_space> + %tile = pto.alloc_tile : !pto.tile_buf + pto.tci ins(%c0_i32 : i32) + outs(%tile : !pto.tile_buf) + pto.tstore ins(%tile : !pto.tile_buf) + outs(%src : memref<1x32xi32, strided<[32, 1], offset: ?>, #pto.address_space>) {layout = #pto.layout, pto.inferred_layout = true} + return + } +} + +// IR: pto.alloc_tile addr = {{.*}} : !pto.tile_buf +// IR: pto.tci ins(%{{.*}}, %{{.*}} : i32, !pto.tile_buf) +// IR-NOT: memref.alloc + +// CPP: TCI<{{.*}}, {{.*}}, int32_t, 0>({{.*}}, {{.*}}, {{.*}}) + +// A5IR: pto.tci ins(%{{.*}} : i32) outs( +// A5IR-NOT: !pto.tile_buf diff --git a/test/lit/pto/tci_tmp_contract_a3_invalid.pto b/test/lit/pto/tci_tmp_contract_a3_invalid.pto new file mode 100644 index 0000000000..1694faf370 --- /dev/null +++ b/test/lit/pto/tci_tmp_contract_a3_invalid.pto @@ -0,0 +1,27 @@ +// RUN: not ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @tci_b32_tmp_too_small() { + %c0_i32 = arith.constant 0 : i32 + %dst = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + // CHECK: error: 'pto.tci' op expects A2/A3 tmp capacity to be at least 768 bytes for 32-bit dst element type + pto.tci ins(%c0_i32, %tmp : i32, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @tci_b16_tmp_too_small() { + %c0_i16 = arith.constant 0 : i16 + %dst = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + // CHECK: error: 'pto.tci' op expects A2/A3 tmp capacity to be at least 1792 bytes for 16-bit dst element type + pto.tci ins(%c0_i16, %tmp : i16, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/pto/tci_ui16_emitc.pto b/test/lit/pto/tci_ui16_emitc.pto index aba69bf2f0..4d435ae994 100644 --- a/test/lit/pto/tci_ui16_emitc.pto +++ b/test/lit/pto/tci_ui16_emitc.pto @@ -11,5 +11,5 @@ module { } } -// A3: TCI<{{.*}}, uint16_t, 0>( -// A3-NOT: TCI<{{.*}}, int16_t, 0>( +// A3: TCI<{{.*}}, {{.*}}, uint16_t, 0>({{.*}}, {{.*}}, {{.*}}) +// A3-NOT: TCI<{{.*}}, {{.*}}, int16_t, 0>( diff --git a/test/lit/pto/tci_ui32_emitc.pto b/test/lit/pto/tci_ui32_emitc.pto index ef6cc2a31c..beb7d241b7 100644 --- a/test/lit/pto/tci_ui32_emitc.pto +++ b/test/lit/pto/tci_ui32_emitc.pto @@ -11,5 +11,5 @@ module { } } -// A3: TCI<{{.*}}, uint32_t, 0>( -// A3-NOT: TCI<{{.*}}, int32_t, 0>( +// A3: TCI<{{.*}}, {{.*}}, uint32_t, 0>({{.*}}, {{.*}}, {{.*}}) +// A3-NOT: TCI<{{.*}}, {{.*}}, int32_t, 0>( diff --git a/test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto b/test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto new file mode 100644 index 0000000000..34f9d32a7c --- /dev/null +++ b/test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto @@ -0,0 +1,24 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --pto-level=level3 %s 2>&1 | FileCheck %s + +module { + func.func @trowexpand_mode1_level3_requires_tmp() { + %addr0 = arith.constant 0 : i64 + %addr1 = arith.constant 1024 : i64 + %addr2 = arith.constant 2048 : i64 + %src0 = pto.alloc_tile addr = %addr0 : !pto.tile_buf + %src1 = pto.alloc_tile addr = %addr1 : !pto.tile_buf + %dst = pto.alloc_tile addr = %addr2 : !pto.tile_buf + // CHECK: error: 'pto.trowexpandadd' op requires explicit tmp for A2/A3 row-expand mode 1 when PlanMemory is skipped + pto.trowexpandadd ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/pto/trowexpand_implicit_tmp_materialization.pto b/test/lit/pto/trowexpand_implicit_tmp_materialization.pto new file mode 100644 index 0000000000..1ffe8abc36 --- /dev/null +++ b/test/lit/pto/trowexpand_implicit_tmp_materialization.pto @@ -0,0 +1,42 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 + +module { + func.func @trowexpand_mode1_add_materializes_tmp() { + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.trowexpandadd ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @trowexpand_mode2_add_keeps_no_tmp() { + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.trowexpandadd ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// A3-LABEL: func.func @trowexpand_mode1_add_materializes_tmp +// A3: pto.alloc_tile addr = {{.*}} : !pto.tile_buf +// A3: pto.trowexpandadd ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) +// A3-LABEL: func.func @trowexpand_mode2_add_keeps_no_tmp +// A3: pto.trowexpandadd ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3-NOT: !pto.tile_buf + +// A5-LABEL: func.func @trowexpand_mode1_add_materializes_tmp +// A5: pto.trowexpandadd ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5-LABEL: func.func @trowexpand_mode2_add_keeps_no_tmp +// A5-NOT: !pto.tile_buf diff --git a/test/lit/pto/trowexpand_tmp_contract_a3_invalid.pto b/test/lit/pto/trowexpand_tmp_contract_a3_invalid.pto new file mode 100644 index 0000000000..bea63f9b29 --- /dev/null +++ b/test/lit/pto/trowexpand_tmp_contract_a3_invalid.pto @@ -0,0 +1,33 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @trowexpand_mode2_rejects_explicit_tmp() { + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + // CHECK: error: 'pto.trowexpandadd' op expects A2/A3 tmp-form trowexpand to use mode 1 + pto.trowexpandadd ins(%src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @trowexpand_mode1_rejects_small_tmp() { + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + // CHECK: error: 'pto.trowexpandmax' op expects A2/A3 trowexpand tmp capacity to be at least 512 bytes, but got 256 bytes + pto.trowexpandmax ins(%src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 8fdaa31afa..7781a4281f 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -3460,6 +3460,9 @@ int mlir::pto::compilePTOASModule( pm.addNestedPass(pto::createPTOFusionRegionGenPass()); } + pm.addNestedPass( + pto::createPTOMaterializeImplicitTmpPass( + effectiveLevel == PTOBuildLevel::Level3)); pm.addNestedPass( pto::createPTORematerializeFixpipeVectorQuantPass()); From 17a59c23f98fda99e79d69ef28f0faa13002ee9e Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 17 Jul 2026 11:28:06 +0800 Subject: [PATCH 044/122] all rest ops --- ...oas-implicit-tmp-materialization-design.md | 295 ++++++ include/PTO/IR/PTOOps.td | 130 +-- lib/PTO/IR/PTO.cpp | 881 +++++++++++++++--- .../InsertSync/InsertSyncAnalysis.cpp | 49 + .../Transforms/PTOMaterializeImplicitTmp.cpp | 685 ++++++++++++++ lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 8 +- lib/PTO/Transforms/PTOToEmitC.cpp | 6 +- test/lit/pto/implicit_tmp_arg_reductions.pto | 31 + .../implicit_tmp_optional_level3_invalid.pto | 24 + ...licit_tmp_optional_ops_materialization.pto | 80 ++ .../implicit_tmp_remaining_level3_invalid.pto | 46 + ...icit_tmp_remaining_ops_materialization.pto | 42 + test/lit/pto/implicit_tmp_row_reductions.pto | 31 + .../pto/implicit_tmp_xor_materialization.pto | 27 + ...ssue533_loop_zero_trip_sync_regression.pto | 8 +- ...533_loop_zero_trip_sync_regression_gss.pto | 8 +- test/lit/pto/issue646_pipev_repeat_prune.pto | 16 +- test/lit/pto/tpow_fp_missing_tmp_invalid.pto | 11 +- 18 files changed, 2143 insertions(+), 235 deletions(-) create mode 100644 test/lit/pto/implicit_tmp_arg_reductions.pto create mode 100644 test/lit/pto/implicit_tmp_optional_level3_invalid.pto create mode 100644 test/lit/pto/implicit_tmp_optional_ops_materialization.pto create mode 100644 test/lit/pto/implicit_tmp_remaining_level3_invalid.pto create mode 100644 test/lit/pto/implicit_tmp_remaining_ops_materialization.pto create mode 100644 test/lit/pto/implicit_tmp_row_reductions.pto create mode 100644 test/lit/pto/implicit_tmp_xor_materialization.pto diff --git a/docs/designs/ptoas-implicit-tmp-materialization-design.md b/docs/designs/ptoas-implicit-tmp-materialization-design.md index 40ebab49f3..719563b370 100644 --- a/docs/designs/ptoas-implicit-tmp-materialization-design.md +++ b/docs/designs/ptoas-implicit-tmp-materialization-design.md @@ -566,6 +566,301 @@ A5 level3 + no tmp 不应因为 tmp 缺失报错。 - A2/A3 动态 `dst.validRow` 且显式 tmp 小于 8192 字节，按保守规则报错。 - A5 显式 tmp 不触发 A2/A3 容量校验。 +## 批量 optional tmp op 分类设计 + +本节按 PTO-ISA tmp 行为对后续待改造 op 做分类。输入列表中的重复项只记录一次： + +```text +TCOLARGMAX, TCOLARGMIN, TROWARGMAX, TROWARGMIN, +TADDDEQRELU, TGATHER, TTRANS, TXOR, TXORS, TPRELU, TMRGSORT, +TROWPROD, TCOLSUM, TROWSUM, TROWMAX, TROWMIN, +TSEL, TSELS, TRSQRT, TPOW, TPOWS, TREM, TREMS, TCVT, +TSORT32, TQUANT +``` + +其中 `TADDDEQRELU` 对应 PTO-ISA 文档 `TAddDeqRelu_zh.md`；当前 PTOAS ODS 中未找到同名 op，先标记为待 IR 接入。 + +### 分类总览 + +| 分类 | Op / 模式 | 设计结论 | +| --- | --- | --- | +| A2/A3 使用 tmp，A5 接受但不使用 | `TCOLARGMAX`、`TCOLARGMIN`、`TROWARGMAX`、`TROWARGMIN`、`TGATHER`、`TROWPROD`、`TROWSUM`、`TROWMAX`、`TROWMIN`、`TSEL`、`TSELS`、`TREM`、`TREMS`、`TQUANT`、`TADDDEQRELU` | level1/2 在 A2/A3 生成真实 scratch；若 A5 C++ 签名仍要求 tmp，则生成不带 MemoryEffects 的 ABI placeholder。 | +| A2/A3 和 A5 都可能使用 tmp | `TCOLSUM(isBinary=true)`、`TSORT32` 非 32 对齐尾部、`TMRGSORT` 多列表归并 format2 | tmp 使用由 op 模式决定，不能只按 arch 判断。 | +| 条件性 tmp，不应无条件 materialize | `TTRANS`、`TCVT`、`TPOW`、`TPOWS`、`TRSQRT`、`TMRGSORT`、`TSORT32` | 需要先判断精度、dtype、layout、format 或尾部条件。 | +| 已从 mandatory tmp 改为 optional tmp | `TTRANS`、`TXOR`、`TXORS`、`TPRELU`、`TROWPROD`、`TROWSUM`、`TROWMAX`、`TROWMIN`、`TROWARGMAX`、`TROWARGMIN`、`TCOLARGMAX`、`TCOLARGMIN`、`TSEL`、`TSELS`、`TREM`、`TREMS` | ODS、parse/print、verifier、MemoryEffects、materialize 和 lowering 已接入。 | +| 当前 PTOAS IR 已有 optional tmp | `TGATHER`、`TCOLSUM`、`TRSQRT`、`TPOW`、`TPOWS`、`TSORT32`、`TQUANT` | 可直接纳入 `pto-materialize-implicit-tmp` 的后续实现。 | +| 当前 PTOAS IR 暂无对应 op | `TADDDEQRELU` | 需先完成 PTOAS IR 接入；`TCVT` 已新增 optional tmp operand。 | + +### 通用规则 + +- level1/level2：只有该 op 在当前 arch / 模式下实际需要 tmp，且 IR 允许省略 tmp 时，才自动 materialize `pto.alloc_tile(no addr)`。 +- level3：若该 op 在当前 arch / 模式下需要 tmp 且用户省略 tmp，则报错；若当前 arch / 模式不使用 tmp，则不强制补 tmp。 +- A5 仅接受但不使用 tmp 的 op：若后端存在 no-tmp overload，则不自动补 tmp；若 C++ 签名仍要求 tmp，则自动生成 ABI placeholder。placeholder 不建模 tmp 的 Read/Write，用户显式 tmp 也不按 A2/A3 容量规则校验。 +- tmp 是 scratch 的 op：MemoryEffects 需要建模为 `Read(tmp) + Write(tmp)`，或在 semantic no-alias side table 中显式加入 `forbidAlias(tmp, dst/output)`。 +- 原 mandatory tmp op 已统一改为 optional，并保持显式 tmp 文本格式兼容。 + +### Arg reduction 类 + +覆盖： + +```text +TCOLARGMAX, TCOLARGMIN, TROWARGMAX, TROWARGMIN +``` + +现状： + +- 这些 op 的 tmp 已改为 optional，并接入 `pto-materialize-implicit-tmp`。 +- A2/A3 使用 tmp；A5 接受 tmp 但不使用。 + +TCOLARGMAX / TCOLARGMIN： + +- tmp dtype 必须与 `src` 一致。 +- tmp 用于索引跟踪和当前比较值临时存储。 +- tmp 容量需要按 `tmpGapEles` 和输出模式计算：当 `srcValidCol >= elemPerRpt` 时，`tmpGapEles = elemPerRpt`；否则 `tmpGapEles = ceil(srcValidCol / elemPerBlock) * elemPerBlock`。 +- half + 纯索引模式是 tmp 使用量最大的组合；其它类型 / 模式下 tmp 中可能只需要区域 0，但自动 materialize 可先采用覆盖最大需求的保守形状。 + +TROWARGMAX / TROWARGMIN： + +- 仅索引模式在 A2/A3 可能不使用 tmp；值+索引模式和两阶段归约需要 tmp。 +- tmp 行数与 `src` 相同；每行 stride 按 PTO-ISA 文档公式计算。 +- 当前 PTOAS ODS 只有单输出索引模式；该模式仍需满足后端显式 tmp 参数签名，因此 level1/2 生成保守同形状 tmp。未来接入值+索引模式后，再按输出模式和归约阶段收紧容量。 + +MemoryEffects / alias： + +- A2/A3 实际使用 tmp 时建模 `Read(src) + Read(tmp) + Write(tmp) + Write(dstIdx/dstVal)`。 +- tmp 不应与同 op 的输出 alias；如果 MemoryEffects 无法覆盖，应加入 `forbidAlias(tmp, dstIdx)` 和必要的 `forbidAlias(tmp, dstVal)`。 + +### Row reduction 类 + +覆盖： + +```text +TROWPROD, TROWSUM, TROWMAX, TROWMIN +``` + +现状： + +- 这些 op 的 tmp 已改为 optional，并接入 `pto-materialize-implicit-tmp`。 +- A2/A3 使用 tmp；A5 接受 tmp 但不使用。 + +tmp 规格： + +- tmp dtype 与 `src` / `dst` 一致。 +- 整数路径最小需要 1 个 vector block：`int32` 为 8 列，`int16` 为 16 列。 +- 浮点路径用于二叉树归约；安全默认形状可设为与 `src` 相同。 +- `TROWPROD` 的安全默认也可设为与 `src` 相同；最小需求为 1 行和 1 个 vector block。 + +Pass 行为： + +- A2/A3 level1/2：若 IR 已支持 optional tmp 且缺省 tmp，则自动生成 vec row-major none-box tmp。 +- A5：生成后端签名需要的 ABI placeholder，但不为其添加 tmp MemoryEffects。 +- level3：A2/A3 需要 tmp 时缺省 tmp 报错；A5 不强制 tmp。 + +MemoryEffects / alias： + +- A2/A3 建模 `Read(src) + Read(tmp) + Write(tmp) + Write(dst)`。 +- tmp 与 `dst` 禁止 alias。 + +### Column sum 类 + +覆盖： + +```text +TCOLSUM +``` + +现状： + +- 当前 PTOAS IR 已支持 optional tmp。 +- no-tmp 形式表示顺序累加，不需要 tmp。 +- `isBinary=true` 时 A2/A3 和 A5 都使用 tmp 做二叉树累加。 + +tmp 规格： + +- tmp dtype 与 `src` / `dst` 一致。 +- tmp 为 vec row-major none-box tile。 +- `tmp.validCol >= src.validCol`。 +- `tmp.validRow >= ceil(src.validRow / 2)`。 + +Pass 行为： + +- 仅当 `isBinary=true` 且缺省 tmp 时自动 materialize。 +- `isBinary=false` 不自动补 tmp。 +- level3 下 `isBinary=true` 且缺省 tmp 报错。 + +MemoryEffects / alias： + +- `isBinary=true` 且 tmp 存在时，建模 `Read(src) + Read(tmp) + Write(tmp) + Write(dst)`。 +- `isBinary=false` 且 tmp 缺省时，保持 no-tmp 顺序累加语义。 + +### Elementwise scratch 类 + +覆盖： + +```text +TXOR, TXORS, TPRELU, TREM, TREMS, TADDDEQRELU, TQUANT +``` + +现状： + +- `TXOR`、`TXORS`、`TPRELU`、`TREM`、`TREMS` 的 tmp 均已改为 optional 并接入 materialize pass。 +- `TQUANT` 当前 PTOAS IR 已支持 optional tmp。 +- `TADDDEQRELU` 当前 PTOAS ODS 中未找到同名 op，需先完成 IR 接入。 + +tmp 规格： + +- `TXOR` / `TXORS`：A2/A3 tmp dtype 与输入输出一致，row-major，容量覆盖 `dst` 有效区域；A5 不使用 tmp。 +- `TPRELU`：A2/A3 tmp dtype 为 `uint8_t`，row-major，`tmp.validRow > dst.validRow`，用于 mask buffer；A5 不使用 tmp。 +- `TREM`：A2/A3 tmp dtype 与 `dst` 一致，至少 2 行和 `dst.validCol` 列；A5 不使用 tmp。 +- `TREMS`：A2/A3 tmp dtype 与 `dst` 一致，至少 1 行和 `dst.validCol` 列；A5 不使用 tmp。 +- `TADDDEQRELU`：A2/A3 tmp dtype 为 `int32_t`，容量至少覆盖 `dst` 有效区域；A5 不使用 tmp。 +- `TQUANT`：A2/A3 tmp 为 FP32，形状与 `src` 同尺寸，用作 FP32 到 S32 转换中间结果；A5 不使用 tmp。 + +Pass 行为： + +- A2/A3 level1/2：缺省 tmp 且 IR 支持 optional tmp 时自动 materialize。 +- A5：`TXOR/TXORS` 因后端签名要求生成 ABI placeholder；其余 op 按各自后端是否存在 no-tmp overload 决定。 +- level3：A2/A3 缺省 tmp 报错；A5 不强制 tmp。 + +MemoryEffects / alias： + +- A2/A3 tmp 是 scratch 时建模 `Read(tmp) + Write(tmp)`。 +- tmp 与 `dst` 禁止 alias；`TXOR/TXORS/TPRELU/TREM/TREMS/TADDDEQRELU` 还应禁止 tmp 与同 op 输入错误 alias。 + +### Mask select 类 + +覆盖： + +```text +TSEL, TSELS +``` + +现状： + +- 当前 PTOAS IR 中 tmp 是 mandatory；隐式 tmp 支持前需要先改成 optional tmp。 +- A2/A3 使用 tmp；A5 接受 tmp 但不使用。 + +tmp 规格： + +- `TSEL`：tmp dtype 为 `uint32_t`，用于 mask buffer。16 位数据类型的 `cmpmaskLen = 4` 个 `uint32_t`；32 位数据类型的 `cmpmaskLen = 2` 个 `uint32_t`。 +- `TSELS`：tmp dtype 与 `src` 一致，至少 1 个元素，用于保存 scalar 和比较 mask。 + +Pass 行为： + +- A2/A3 level1/2：缺省 tmp 时自动 materialize。 +- A5：不自动补 tmp。 +- level3：A2/A3 缺省 tmp 报错；A5 不强制 tmp。 + +MemoryEffects / alias： + +- A2/A3 建模 `Read(mask/src) + Read(tmp) + Write(tmp) + Write(dst)`。 +- tmp 与 `dst` 禁止 alias。 + +### Data movement / layout 类 + +覆盖： + +```text +TGATHER, TTRANS, TCVT +``` + +TGATHER： + +- 当前 PTOAS IR 已支持 optional tmp。 +- index form：A2/A3 C++ API 需要 tmp；tmp dtype 与 indices dtype 一致，shape 覆盖 indices；A5 不使用 tmp。 +- compare form：A2/A3 tmp 是合并暂存缓冲区，包含 `cmpsTmp`、`indexTmp`、`cvtTmp` 三个区域；最小字节数按 PTO-ISA 文档公式计算；A5 不使用 tmp。 +- mask form 不使用 tmp，不应自动补 tmp。 +- A2/A3 level3 下 index / compare form 缺省 tmp 报错。 + +TTRANS： + +- 当前 PTOAS IR 中 tmp 已改为 optional，并接入 materialize pass。 +- tmp 只在满足高效转置路径条件时使用；scalar copy 和部分 layout 转换不需要 tmp。 +- 静态满足 stride 条件时生成与 src 同形状的保守 scratch；不满足并走 scalar copy 时生成 32 字节 ABI placeholder。 +- 只有真实 scratch 进入 `Read(tmp) + Write(tmp)` MemoryEffects；scalar-copy placeholder 不建模内存访问。 + +TCVT： + +- 当前 PTOAS IR 已新增 optional tmp operand，并完成 parser/printer、verifier、MemoryEffects、materialize 和 EmitC lowering。 +- A2/A3 仅在 `SaturationMode::OFF` 的 PyTorch 兼容非饱和窄化路径使用 tmp：`float -> int16`、`half -> int16`、`half -> int8`。 +- 其它转换不需要 tmp，不应自动 materialize。 +- tmp 按字节规划，容量使用 PTO-ISA 的 `tmpFloatToInt16Bytes`、`tmpHalfToInt16Bytes`、`tmpHalfToInt8Bytes` 公式。 +- level3 下上述三种路径缺省 tmp 报错；其它转换继续使用 no-tmp overload。 + +MemoryEffects / alias： + +- 只有实际使用 tmp 的路径建模 `Read(tmp) + Write(tmp)`。 +- tmp 与 `dst` 禁止 alias。 + +### Sort / merge 类 + +覆盖： + +```text +TSORT32, TMRGSORT +``` + +TSORT32： + +- 当前 PTOAS IR 已支持 optional tmp。 +- 3 参数形式适用于 `validCol` 已按 32 对齐的路径，不需要 tmp。 +- 4 参数形式用于非 32 对齐尾部，通过 tmp 保存填充后的行或尾块副本。 +- tmp dtype 与 `src` 一致；容量按 PTO-ISA `tmpSize` 公式计算，不能固定为 8KB。 +- level1/2 仅在能静态证明存在非 32 对齐尾部时自动补 tmp；否则保持 no-tmp。 +- level3 下非 32 对齐尾部缺省 tmp 报错。 + +TMRGSORT： + +- 当前 PTOAS IR 已支持 optional tmp；缺省 format2 使用显式 `no_tmp` 中间语法消除 src/tmp 个数歧义。 +- format1 单输入 block sort 不需要 tmp。 +- format2 多列表归并需要 tmp 和 executed list。 +- level1/2 对 format2 `no_tmp` 生成一行 row-major tmp，`tmp.cols = sum(src.cols)`；format1 不补 tmp。 +- level3 下 format2 `no_tmp` 报错。 + +MemoryEffects / alias： + +- 使用 tmp 的 sort / merge 路径建模 `Read(tmp) + Write(tmp)`。 +- tmp 与 `dst` 禁止 alias；`TMRGSORT` 还需考虑 `executed` output 的写 effect。 + +### Pow / rsqrt 类 + +覆盖： + +```text +TPOW, TPOWS, TRSQRT +``` + +TPOW / TPOWS： + +- 当前 PTOAS IR 已支持 optional tmp。 +- A2/A3 浮点路径使用 tmp；整数路径不使用 tmp。 +- A5 接受 tmp 但不使用。 +- tmp dtype 与 `dst` / `base` 一致，容量覆盖 `dst` 有效区域。 +- level1/2：A2/A3 浮点路径缺省 tmp 时自动 materialize；整数路径不补 tmp。 +- level3：A2/A3 浮点路径缺省 tmp 报错；整数路径不强制 tmp。 + +TRSQRT： + +- 当前 PTOAS IR 已支持 optional tmp。 +- no-tmp 默认实现不需要 tmp。 +- 带 tmp overload 当前仅作为 API 兼容 / 未来高精度路径保留，现阶段不自动 materialize。 +- 用户显式 tmp 时，A5 不按 A2/A3 scratch 容量规则校验。 + +MemoryEffects / alias： + +- `TPOW/TPOWS` 只有浮点 tmp-backed 路径建模 `Read(tmp) + Write(tmp)`。 +- `TRSQRT` 现阶段不因缺省 tmp 增加 MemoryEffects。 + +### 批量 op 测试策略 + +后续按类别实现时，每一类至少补以下 lit： + +- 自动补 tmp：level1/2 + A2/A3 + 缺省 tmp，检查 `pto.alloc_tile(no addr)` 经 memplan 后带 `addr`。 +- A5 不补 tmp：A5 + 缺省 tmp，检查保持 no-tmp 形态。 +- level3 负例：A2/A3 + 需要 tmp + 缺省 tmp 报错。 +- verifier 负例：显式 tmp dtype / shape / capacity / layout 不满足对应 op 规格时报错。 +- EmitC overload：需要 tmp 的路径最终走 tmp-aware C++ 调用；不需要 tmp 的路径不强行切换。 + ## 后续扩展 后续新增其它 optional tmp op 时，需要补充一个 op-specific 小节，并明确： diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 69f1e6c149..aa43df7569 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -861,16 +861,12 @@ def TTransOp : PTO_TOp<"ttrans", [ }]; let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); let results = (outs); - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let hasVerifier = 1; @@ -4070,7 +4066,7 @@ def TColArgMaxOp : PTO_TOp<"tcolargmax", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -4083,11 +4079,7 @@ def TColArgMaxOp : PTO_TOp<"tcolargmax", [ ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } }]; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; } def TColMinOp : PTO_TOp<"tcolmin", [ @@ -4127,7 +4119,7 @@ def TColArgMinOp : PTO_TOp<"tcolargmin", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -4140,11 +4132,7 @@ def TColArgMinOp : PTO_TOp<"tcolargmin", [ ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } }]; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; } def TColSumOp : PTO_TOp<"tcolsum", [ @@ -4214,6 +4202,7 @@ def TCvtOp : PTO_TOp<"tcvt", [ let arguments = (ins PTODpsType:$src, + Optional:$tmp, PTODpsType:$dst, DefaultValuedAttr:$rmode, DefaultValuedAttr:$sat_mode @@ -5149,6 +5138,7 @@ def TMrgSortOp: PTO_TOp<"tmrgsort", [ let extraClassDeclaration = [{ bool isFormat1() { return getSrcs().size() == 1u && getBlockLen() && getDsts().size() == 1u; } bool isFormat2() { return getSrcs().size() >= 2u && getSrcs().size() <= 4u && getTmp() && getDsts().size() == 1u && getExcuted(); } + bool isFormat2WithoutTmp() { return getSrcs().size() >= 2u && getSrcs().size() <= 4u && !getTmp() && !getBlockLen() && getDsts().size() == 1u && getExcuted(); } Value getSrc() { return getSrcs().front(); } Value getDst() { return getDsts().front(); } ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstsMutable(); } @@ -5565,7 +5555,7 @@ def TPReluOp: PTO_TOp<"tprelu", [ let arguments = (ins PTODpsType:$src0, PTODpsType:$src1, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -5573,11 +5563,7 @@ def TPReluOp: PTO_TOp<"tprelu", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src0 `,` $src1 `,` $tmp `:` qualified(type($src0)) `,` qualified(type($src1)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -5815,7 +5801,7 @@ def TRemOp: PTO_TOp<"trem", [ let arguments = (ins PTODpsType:$src0, PTODpsType:$src1, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst, DefaultValuedAttr:$precisionType ); @@ -5824,11 +5810,7 @@ def TRemOp: PTO_TOp<"trem", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src0 `,` $src1 `,` $tmp `:` qualified(type($src0)) `,` qualified(type($src1)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -5849,7 +5831,7 @@ def TRemSOp: PTO_TOp<"trems", [ let arguments = (ins PTODpsType:$src, ScalarType:$scalar, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -5857,11 +5839,7 @@ def TRemSOp: PTO_TOp<"trems", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $scalar `,` $tmp `:` qualified(type($src)) `,` type($scalar) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6154,7 +6132,7 @@ def TRowMaxOp: PTO_TOp<"trowmax", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6162,11 +6140,7 @@ def TRowMaxOp: PTO_TOp<"trowmax", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6183,7 +6157,7 @@ def TRowArgMaxOp: PTO_TOp<"trowargmax", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6191,11 +6165,7 @@ def TRowArgMaxOp: PTO_TOp<"trowargmax", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6215,7 +6185,7 @@ def TRowMinOp: PTO_TOp<"trowmin", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6223,11 +6193,7 @@ def TRowMinOp: PTO_TOp<"trowmin", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6244,7 +6210,7 @@ def TRowArgMinOp: PTO_TOp<"trowargmin", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6252,11 +6218,7 @@ def TRowArgMinOp: PTO_TOp<"trowargmin", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6276,7 +6238,7 @@ def TRowSumOp: PTO_TOp<"trowsum", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6284,11 +6246,7 @@ def TRowSumOp: PTO_TOp<"trowsum", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6305,7 +6263,7 @@ def TRowProdOp: PTO_TOp<"trowprod", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6313,11 +6271,7 @@ def TRowProdOp: PTO_TOp<"trowprod", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6433,7 +6387,7 @@ def TSelOp: PTO_TOp<"tsel", [ PTODpsType:$mask, PTODpsType:$src0, PTODpsType:$src1, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6441,11 +6395,7 @@ def TSelOp: PTO_TOp<"tsel", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $mask `,` $src0 `,` $src1 `,` $tmp `:` qualified(type($mask)) `,` qualified(type($src0)) `,` qualified(type($src1)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6469,7 +6419,7 @@ def TSelSOp: PTO_TOp<"tsels", [ let arguments = (ins PTODpsType:$mask, PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, ScalarType:$scalar, PTODpsType:$dst ); @@ -6478,11 +6428,7 @@ def TSelSOp: PTO_TOp<"tsels", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $mask `,` $src `,` $tmp `,` $scalar `:` qualified(type($mask)) `,` qualified(type($src)) `,` qualified(type($tmp)) `,` type($scalar) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6857,7 +6803,7 @@ def TXorSOp: PTO_TOp<"txors", [ let arguments = (ins PTODpsType:$src, AnySignlessInteger:$scalar, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6865,11 +6811,7 @@ def TXorSOp: PTO_TOp<"txors", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $scalar `,` $tmp `:` qualified(type($src)) `,` type($scalar) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6896,7 +6838,7 @@ def TXorOp: PTO_TOp<"txor", [ let arguments = (ins PTODpsType:$src0, PTODpsType:$src1, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6904,11 +6846,7 @@ def TXorOp: PTO_TOp<"txor", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src0 `,` $src1 `,` $tmp `:` qualified(type($src0)) `,` qualified(type($src1)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 90362862ab..984e0ad7d4 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -2238,6 +2238,25 @@ static LogicalResult verifyTColArgReductionOpA2A3(Operation *op, Type srcTy, return success(); } +static LogicalResult verifyTColArgReductionNoTmp(Operation *op, Type srcTy, + Type dstTy) { + if (failed(verifyNDStyleVecTile(op, srcTy, "src")) || + failed(verifyColArgReductionDstLayout(op, dstTy, "dst")) || + failed(verifyColReductionValidRegion(op, srcTy, dstTy, + /*requireNonZeroSrc=*/true))) + return failure(); + Type srcElemTy = getElemTy(srcTy); + unsigned srcElemBits = srcElemTy ? getPTOStorageElemBitWidth(srcElemTy) : 0; + if (!(mlir::isa(srcElemTy) && + (srcElemBits == 8 || srcElemBits == 16 || srcElemBits == 32))) + return op->emitOpError( + "expects src element type to be 1, 2, or 4 bytes wide"); + auto dstInt = dyn_cast(getElemTy(dstTy)); + if (!dstInt || dstInt.getWidth() != 32) + return op->emitOpError("expects dst element type to be i32 or ui32"); + return success(); +} + static LogicalResult verifyTColArgReductionOpA5(Operation *op, Type srcTy, Type tmpTy, Type dstTy) { if (failed(verifyNDStyleVecTile(op, srcTy, "src")) || @@ -2371,6 +2390,22 @@ static LogicalResult verifyTRowArgReductionOpA2A3(Operation *op, Type srcTy, return success(); } +static LogicalResult verifyTRowArgReductionNoTmp(Operation *op, Type srcTy, + Type dstTy) { + if (failed(verifyRowReductionSrcLayout(op, srcTy, "src")) || + failed(verifyRowReductionDstLayout(op, dstTy, "dst")) || + failed(verifyRowReductionValidRegion(op, srcTy, dstTy, + /*allowEmptyMarker=*/false))) + return failure(); + Type srcElem = getElemTy(srcTy); + if (!isSupportedRowReductionElemType(srcElem)) + return op->emitOpError("expects src element type to be i16/i32/f16/f32"); + auto dstInt = dyn_cast(getElemTy(dstTy)); + if (!dstInt || dstInt.getWidth() != 32) + return op->emitOpError("expects dst element type to be i32 or ui32"); + return success(); +} + static LogicalResult verifyTRowArgReductionOpA5(Operation *op, Type srcTy, Type tmpTy, Type dstTy) { if (failed(verifyRowReductionSrcLayout(op, srcTy, "src")) || @@ -6194,6 +6229,9 @@ LogicalResult pto::TColMaxOp::verify() { } LogicalResult pto::TColArgMaxOp::verify() { + if (!getTmp()) + return verifyTColArgReductionNoTmp(getOperation(), getSrc().getType(), + getDst().getType()); auto verifyA2A3 = [&]() -> LogicalResult { return verifyTColArgReductionOpA2A3(*this, getSrc().getType(), getTmp().getType(), getDst().getType()); @@ -6216,6 +6254,9 @@ LogicalResult pto::TColMinOp::verify() { } LogicalResult pto::TColArgMinOp::verify() { + if (!getTmp()) + return verifyTColArgReductionNoTmp(getOperation(), getSrc().getType(), + getDst().getType()); auto verifyA2A3 = [&]() -> LogicalResult { return verifyTColArgReductionOpA2A3(*this, getSrc().getType(), getTmp().getType(), getDst().getType()); @@ -6325,11 +6366,8 @@ LogicalResult pto::TColSumOp::verify() { return failure(); bool hasTmp = (bool)getTmp(); bool hasIsBinary = (bool)getIsBinaryAttr(); - if (hasTmp != hasIsBinary) { - if (hasTmp) - return emitOpError("tmp operand requires isBinary attribute"); - return emitOpError("isBinary attribute requires tmp operand"); - } + if (hasTmp && !hasIsBinary) + return emitOpError("tmp operand requires isBinary attribute"); if (getTmp()) { Type tmpTy = getTmp().getType(); if (failed(verifyNDStyleVecTile(*this, tmpTy, "tmp"))) @@ -6357,11 +6395,8 @@ LogicalResult pto::TColSumOp::verify() { return failure(); bool hasTmp = (bool)getTmp(); bool hasIsBinary = (bool)getIsBinaryAttr(); - if (hasTmp != hasIsBinary) { - if (hasTmp) - return emitOpError("tmp operand requires isBinary attribute"); - return emitOpError("isBinary attribute requires tmp operand"); - } + if (hasTmp && !hasIsBinary) + return emitOpError("tmp operand requires isBinary attribute"); if (getTmp()) { Type tmpTy = getTmp().getType(); if (failed(verifyNDStyleVecTile(*this, tmpTy, "tmp"))) @@ -6419,14 +6454,65 @@ llvm::LogicalResult mlir::pto::TCvtOp::verify() { return failure(); Type srcElem = getElemTy(srcTy); Type dstElem = getElemTy(dstTy); + auto needsTmp = [&]() { + if (getSatMode() != pto::SaturationMode::OFF) + return false; + return (srcElem.isF32() && dstElem.isInteger(16)) || + (srcElem.isF16() && + (dstElem.isInteger(16) || dstElem.isInteger(8))); + }; + auto verifyTmp = [&]() -> LogicalResult { + if (!getTmp()) + return success(); + Type tmpTy = getTmp().getType(); + if (failed(verifyVecTileCommon(*this, tmpTy, "tmp"))) + return failure(); + if (!needsTmp()) + return success(); + auto srcShape = getShapeVec(srcTy); + auto dstValid = getValidShapeVec(dstTy); + if (srcShape.size() != 2 || dstValid.size() != 2 || + llvm::is_contained(srcShape, ShapedType::kDynamic) || + llvm::is_contained(dstValid, ShapedType::kDynamic)) + return emitOpError( + "expects static src shape and dst valid_shape to verify tcvt tmp"); + int64_t rows = dstValid[0], cols = dstValid[1]; + int64_t requiredBytes = 0; + if (rows > 0 && cols > 0 && srcElem.isF32()) { + int64_t head = 4 * 64 * std::min(cols / 64, 255); + int64_t remainder = cols % 64; + int64_t tail = remainder == 0 + ? 0 + : 32 * ((std::min(rows, 255) - 1) * + (srcShape[1] / 8) + + llvm::divideCeil(remainder, int64_t{8})); + requiredBytes = std::max(head, tail); + } else if (cols > 0 && srcElem.isF16()) { + int64_t width = std::min(cols, 64); + int64_t halfToI16 = 32 * llvm::divideCeil(width, int64_t{8}); + int64_t halfToI8 = std::max( + halfToI16, + 128 + 32 * static_cast( + llvm::divideCeil(width, int64_t{16}))); + requiredBytes = dstElem.isInteger(8) ? halfToI8 : halfToI16; + } + auto tmpBytes = getStaticByteSize(tmpTy); + if (!tmpBytes || *tmpBytes < static_cast(requiredBytes)) + return emitOpError() + << "expects tcvt tmp capacity to be at least " << requiredBytes + << " bytes"; + return success(); + }; auto verifyA2A3 = [&]() -> LogicalResult { if (isPTOLowPrecisionType(srcElem) || isPTOLowPrecisionType(dstElem)) return emitOpError("expects A2/A3 tcvt low-precision element types to be unsupported"); - return success(); + return verifyTmp(); }; auto verifyA5 = [&]() -> LogicalResult { if (!isA5SupportedTCvtPair(srcElem, dstElem)) return emitOpError("expects A5 tcvt low-precision type pairs to match PTO-ISA support"); + if (getTmp() && failed(verifyVecTileCommon(*this, getTmp().getType(), "tmp"))) + return failure(); return success(); }; return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); @@ -7759,14 +7845,18 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() { if (getAxisAttr()) return emitOpError("axis attribute must not be provided without maskPattern"); if (getCdst() || getKValue()) { - if (!getCdst() || !getKValue() || !getTmp()) - return emitOpError("compare-form tgather expects dst, cdst, kValue, and tmp"); + if (!getCdst() || !getKValue()) + return emitOpError("compare-form tgather expects dst, cdst, and kValue"); if (getIndices()) return emitOpError("compare-form tgather does not take indices"); + if (!getTmp()) + return success(); return verifyCompareForm(/*allowA5SrcTypes=*/false); } - if (!getIndices() || !getTmp()) - return emitOpError("index-form tgather expects both indices and tmp"); + if (!getIndices()) + return emitOpError("index-form tgather expects indices"); + if (!getTmp()) + return success(); return verifyIndexForm(/*allow16BitIndices=*/false, /*allowA5ElemTypes=*/false); }; @@ -7779,14 +7869,18 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() { if (getAxisAttr()) return emitOpError("axis attribute must not be provided without maskPattern"); if (getCdst() || getKValue()) { - if (!getCdst() || !getKValue() || !getTmp()) - return emitOpError("compare-form tgather expects dst, cdst, kValue, and tmp"); + if (!getCdst() || !getKValue()) + return emitOpError("compare-form tgather expects dst, cdst, and kValue"); if (getIndices()) return emitOpError("compare-form tgather does not take indices"); + if (!getTmp()) + return success(); return verifyCompareForm(/*allowA5SrcTypes=*/true); } if (!getIndices()) return emitOpError("index-form tgather expects indices"); + if (!getTmp()) + return success(); return verifyIndexForm(/*allow16BitIndices=*/true, /*allowA5ElemTypes=*/true); }; @@ -9540,6 +9634,8 @@ LogicalResult MGatherOp::verify() { void mlir::pto::TCvtOp::print(OpAsmPrinter &p) { p << " ins(" << getSrc(); + if (getTmp()) + p << ", " << getTmp(); Builder builder(getContext()); NamedAttrList attrs; for (auto attr : (*this)->getAttrs()) { @@ -9551,18 +9647,28 @@ void mlir::pto::TCvtOp::print(OpAsmPrinter &p) { } p.printOptionalAttrDict(attrs.getAttrs()); p << " : " << getSrc().getType(); + if (getTmp()) + p << ", " << getTmp().getType(); p << ") outs(" << getDst() << " : " << getDst().getType() << ")"; } ParseResult mlir::pto::TCvtOp::parse(OpAsmParser &parser, OperationState &result) { - OpAsmParser::UnresolvedOperand src, dst; - Type srcTy, dstTy; + OpAsmParser::UnresolvedOperand src, tmp, dst; + Type srcTy, tmpTy, dstTy; + bool hasTmp = false; if (parser.parseKeyword("ins") || parser.parseLParen() || parser.parseOperand(src)) return failure(); + if (succeeded(parser.parseOptionalComma())) { + if (parser.parseOperand(tmp)) + return failure(); + hasTmp = true; + } NamedAttrList attrs; if (parser.parseOptionalAttrDict(attrs) || parser.parseColonType(srcTy)) return failure(); + if (hasTmp && (parser.parseComma() || parser.parseType(tmpTy))) + return failure(); if (auto satmode = attrs.get("satmode")) { attrs.erase("satmode"); if (attrs.get("sat_mode")) @@ -9576,8 +9682,12 @@ ParseResult mlir::pto::TCvtOp::parse(OpAsmParser &parser, OperationState &result return failure(); if (parser.resolveOperand(src, srcTy, result.operands) || + (hasTmp && parser.resolveOperand(tmp, tmpTy, result.operands)) || parser.resolveOperand(dst, dstTy, result.operands)) return failure(); + result.addAttribute( + "operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr({1, hasTmp ? 1 : 0, 1})); return success(); } @@ -9586,13 +9696,17 @@ void mlir::pto::TMrgSortOp::print(OpAsmPrinter &p) { p << " ins(" << getSrc() << ", " << getBlockLen() << " : " << getSrc().getType() << ", " << getBlockLen().getType() << ") outs(" << getDst() << " : " << getDst().getType() << ")"; - } else if (isFormat2()) { + } else if (isFormat2() || isFormat2WithoutTmp()) { p << " ins("; llvm::interleaveComma(getSrcs(), p, [&](Value src) { p << src; }); - p << ", " << getTmp(); + if (getTmp()) + p << ", " << getTmp(); + else + p << " no_tmp"; p << " {exhausted = " << (getExhausted() ? "true" : "false") << "} : "; llvm::interleaveComma(getSrcs().getTypes(), p, [&](Type ty) { p << ty; }); - p << ", " << getTmp().getType(); + if (getTmp()) + p << ", " << getTmp().getType(); p << ") outs(" << getDst() << ", " << getExcuted() << " : " << getDst().getType() << ", " << getExcuted().getType() << ")"; } else { @@ -9637,10 +9751,15 @@ ParseResult mlir::pto::TMrgSortOp::parse(OpAsmParser &parser, OperationState &re return failure(); srcs.push_back(next); } - if (srcs.size() < 3 || srcs.size() > 5) - return parser.emitError(parser.getCurrentLocation(), - "tmrgsort format2 expects 2 to 4 src operands plus one tmp operand"); - OpAsmParser::UnresolvedOperand tmpOp = srcs.pop_back_val(); + bool noTmp = succeeded(parser.parseOptionalKeyword("no_tmp")); + if ((noTmp && (srcs.size() < 2 || srcs.size() > 4)) || + (!noTmp && (srcs.size() < 3 || srcs.size() > 5))) + return parser.emitError( + parser.getCurrentLocation(), + "tmrgsort format2 expects 2 to 4 src operands and optional no_tmp marker"); + OpAsmParser::UnresolvedOperand tmpOp; + if (!noTmp) + tmpOp = srcs.pop_back_val(); bool exhaustedVal = false; if (parser.parseOptionalLBrace().succeeded()) { if (parser.parseKeyword("exhausted") || parser.parseEqual()) @@ -9664,10 +9783,13 @@ ParseResult mlir::pto::TMrgSortOp::parse(OpAsmParser &parser, OperationState &re return failure(); srcTypes.push_back(nextTy); } - if (srcTypes.size() != srcs.size() + 1 || parser.parseRParen() || + if (srcTypes.size() != srcs.size() + (noTmp ? 0 : 1) || + parser.parseRParen() || parser.parseKeyword("outs") || parser.parseLParen()) return failure(); - Type tmpTy = srcTypes.pop_back_val(); + Type tmpTy; + if (!noTmp) + tmpTy = srcTypes.pop_back_val(); OpAsmParser::UnresolvedOperand dstOp, excutedOp; Type dstTy, excutedTy; if (parser.parseOperand(dstOp) || parser.parseComma() || parser.parseOperand(excutedOp) || @@ -9676,10 +9798,11 @@ ParseResult mlir::pto::TMrgSortOp::parse(OpAsmParser &parser, OperationState &re return failure(); result.addAttribute("operandSegmentSizes", parser.getBuilder().getDenseI32ArrayAttr( - {static_cast(srcs.size()), 0, 1, 1, 1})); + {static_cast(srcs.size()), 0, 1, + noTmp ? 0 : 1, 1})); if (parser.resolveOperands(srcs, srcTypes, parser.getCurrentLocation(), result.operands) || parser.resolveOperand(dstOp, dstTy, result.operands) || - parser.resolveOperand(tmpOp, tmpTy, result.operands) || + (!noTmp && parser.resolveOperand(tmpOp, tmpTy, result.operands)) || parser.resolveOperand(excutedOp, excutedTy, result.operands)) return failure(); if (parser.parseOptionalAttrDict(result.attributes)) @@ -9720,36 +9843,40 @@ mlir::LogicalResult mlir::pto::TMrgSortOp::verify() { } return mlir::success(); } - if (isFormat2()) { + if (isFormat2() || isFormat2WithoutTmp()) { for (Value v : getSrcs()) if (!isPTOShapedLike(v.getType())) return emitOpError() << "format2 expects PTO shaped-like type for each src"; if (getSrcs().size() < 2u || getSrcs().size() > 4u) return emitOpError() << "format2 expects 2 to 4 srcs"; - if (getDsts().size() != 1u || !getTmp() || !getExcuted()) - return emitOpError() << "format2 expects ins(srcs..., tmp), outs(dst), and excuted=vector"; + if (getDsts().size() != 1u || !getExcuted()) + return emitOpError() + << "format2 expects 2 to 4 srcs, one dst, and excuted=vector"; Type dstTy = getDst().getType(); - Type tmpTy = getTmp().getType(); - if (!isPTOShapedLike(dstTy) || !isPTOShapedLike(tmpTy)) + Type tmpTy = getTmp() ? getTmp().getType() : Type{}; + if (!isPTOShapedLike(dstTy) || + (tmpTy && !isPTOShapedLike(tmpTy))) return emitOpError() << "format2 dst/tmp must be PTO shaped-like"; auto excutedTy = mlir::dyn_cast(getExcuted().getType()); if (!excutedTy || excutedTy.getRank() != 1 || excutedTy.getNumElements() != 4 || !excutedTy.getElementType().isInteger(16)) return emitOpError() << "format2 excuted must be vector<4xi16>"; Type elemTy = getElemTy(dstTy); - if (elemTy != getElemTy(tmpTy)) + if (tmpTy && elemTy != getElemTy(tmpTy)) return emitOpError() << "format2 expects dst/tmp element types to match"; auto dstShape = getShapeVec(dstTy); - auto tmpShape = getShapeVec(tmpTy); - if (dstShape.size() != 2 || tmpShape.size() != 2) + auto tmpShape = tmpTy ? getShapeVec(tmpTy) : SmallVector{}; + if (dstShape.size() != 2 || (tmpTy && tmpShape.size() != 2)) return emitOpError() << "format2 expects dst/tmp to be rank-2 tile-shaped"; if ((dstShape[0] != mlir::ShapedType::kDynamic && dstShape[0] != 1) || - (tmpShape[0] != mlir::ShapedType::kDynamic && tmpShape[0] != 1)) + (tmpTy && tmpShape[0] != mlir::ShapedType::kDynamic && + tmpShape[0] != 1)) return emitOpError() << "format2 expects dst/tmp rows == 1"; - if (dstShape[1] != mlir::ShapedType::kDynamic && + if (tmpTy && dstShape[1] != mlir::ShapedType::kDynamic && tmpShape[1] != mlir::ShapedType::kDynamic && tmpShape[1] < dstShape[1]) return emitOpError() << "format2 expects tmp.cols >= dst.cols"; + int64_t requiredTmpCols = 0; for (Value src : getSrcs()) { Type srcTy = src.getType(); auto srcShape = getShapeVec(srcTy); @@ -9759,7 +9886,17 @@ mlir::LogicalResult mlir::pto::TMrgSortOp::verify() { return emitOpError() << "format2 expects src rows == 1"; if (getElemTy(srcTy) != elemTy) return emitOpError() << "format2 expects src/dst/tmp element types to match"; + if (srcShape[1] == mlir::ShapedType::kDynamic) + requiredTmpCols = mlir::ShapedType::kDynamic; + else if (requiredTmpCols != mlir::ShapedType::kDynamic) + requiredTmpCols += srcShape[1]; } + if (tmpTy && requiredTmpCols != mlir::ShapedType::kDynamic && + tmpShape[1] != mlir::ShapedType::kDynamic && + tmpShape[1] < requiredTmpCols) + return emitOpError() + << "format2 expects tmp.cols >= sum(src.cols) = " + << requiredTmpCols; return mlir::success(); } return emitOpError() << "tmrgsort expects format1 (1 src + blockLen + 1 dst) or " @@ -10275,16 +10412,17 @@ mlir::LogicalResult mlir::pto::TPReluOp::verify() { auto verifyCommon = [&]() -> FailureOr> { Type t0 = getSrc0().getType(); Type t1 = getSrc1().getType(); - Type tt = getTmp().getType(); + Type tt = getTmp() ? getTmp().getType() : Type{}; Type td = getDst().getType(); if (failed(verifyTileBufCommon(*this, t0, "src0")) || failed(verifyTileBufCommon(*this, t1, "src1")) || - failed(verifyTileBufCommon(*this, tt, "tmp")) || failed(verifyTileBufCommon(*this, td, "dst"))) return failure(); + if (tt && failed(verifyTileBufCommon(*this, tt, "tmp"))) + return failure(); - Type e0 = getElemTy(t0), e1 = getElemTy(t1), et = getElemTy(tt), ed = getElemTy(td); - if (!e0 || !e1 || !et || !ed) { + Type e0 = getElemTy(t0), e1 = getElemTy(t1), ed = getElemTy(td); + if (!e0 || !e1 || !ed) { emitOpError("failed to get element type for operands"); return failure(); } @@ -10317,6 +10455,8 @@ mlir::LogicalResult mlir::pto::TPReluOp::verify() { if (failed(tysOr)) return failure(); auto [t0, t1, tt, td] = *tysOr; + if (!tt) + return success(); Type tmpElem = getElemTy(tt); auto tmpIntTy = mlir::dyn_cast(tmpElem); if (!tmpIntTy || tmpIntTy.getWidth() != 8) @@ -10358,7 +10498,7 @@ mlir::LogicalResult mlir::pto::TPReluOp::verify() { (void)t0; (void)t1; (void)td; - if (failed(verifyVecTileCommon(*this, tt, "tmp"))) + if (tt && failed(verifyVecTileCommon(*this, tt, "tmp"))) return failure(); return success(); }; @@ -10859,11 +10999,9 @@ mlir::LogicalResult mlir::pto::TRemOp::verify() { Type src0Ty = getSrc0().getType(); Type src1Ty = getSrc1().getType(); - Type tmpTy = getTmp().getType(); Type dstTy = getDst().getType(); if (failed(verifyTileBufCommon(*this, src0Ty, "src0")) || failed(verifyTileBufCommon(*this, src1Ty, "src1")) || - failed(verifyTileBufCommon(*this, tmpTy, "tmp")) || failed(verifyTileBufCommon(*this, dstTy, "dst"))) return failure(); if (failed(verifyTileBufSameElemType(*this, src0Ty, src1Ty, "src0", "src1")) || @@ -10875,11 +11013,30 @@ mlir::LogicalResult mlir::pto::TRemOp::verify() { !isRowMajorTileBuf(dstTy)) return emitOpError("expects src0, src1, and dst to use row-major layout"); auto dstValid = getValidShapeVec(dstTy); + + Type elem = getElemTy(src0Ty); + if (!getTmp()) { + auto verifyA2A3NoTmp = [&]() -> LogicalResult { + if (!(elem.isInteger(32) || elem.isF32())) + return emitOpError("expects A2/A3 trem element type to be i32/f32"); + return success(); + }; + auto verifyA5NoTmp = [&]() -> LogicalResult { + if (!(elem.isInteger(32) || elem.isInteger(16) || elem.isF16() || + elem.isF32())) + return emitOpError( + "expects A5 trem element type to be i32/i16/f16/f32"); + return success(); + }; + return dispatchVerifierByArch(getOperation(), verifyA2A3NoTmp, + verifyA5NoTmp); + } + Type tmpTy = getTmp().getType(); + if (failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) + return failure(); auto tmpValid = getValidShapeVec(tmpTy); if (dstValid.size() != 2 || tmpValid.size() != 2) return emitOpError("expects tmp and dst to be rank-2 tiles"); - - Type elem = getElemTy(src0Ty); auto verifyA2A3 = [&]() -> LogicalResult { if (failed(verifyVecTileCommon(*this, tmpTy, "tmp"))) return failure(); @@ -10914,11 +11071,9 @@ mlir::LogicalResult mlir::pto::TFModOp::verify() { mlir::LogicalResult mlir::pto::TRemSOp::verify() { Type ts = getSrc().getType(); - Type tt = getTmp().getType(); Type td = getDst().getType(); Type scalarTy = getScalar().getType(); if (failed(verifyTileBufCommon(*this, ts, "src")) || - failed(verifyTileBufCommon(*this, tt, "tmp")) || failed(verifyTileBufCommon(*this, td, "dst"))) return failure(); if (failed(verifyTileBufSameElemType(*this, ts, td, "src", "dst")) || @@ -10930,6 +11085,25 @@ mlir::LogicalResult mlir::pto::TRemSOp::verify() { if (scalarTy != elem) return emitOpError("expects scalar type to match the tile element type"); auto dstValid = getValidShapeVec(td); + if (!getTmp()) { + auto verifyA2A3NoTmp = [&]() -> LogicalResult { + if (!(elem.isInteger(32) || elem.isF32())) + return emitOpError("expects A2/A3 trems element type to be i32/f32"); + return success(); + }; + auto verifyA5NoTmp = [&]() -> LogicalResult { + if (!(elem.isInteger(32) || elem.isInteger(16) || elem.isF16() || + elem.isF32())) + return emitOpError( + "expects A5 trems element type to be i32/i16/f16/f32"); + return success(); + }; + return dispatchVerifierByArch(getOperation(), verifyA2A3NoTmp, + verifyA5NoTmp); + } + Type tt = getTmp().getType(); + if (failed(verifyTileBufCommon(*this, tt, "tmp"))) + return failure(); auto tmpValid = getValidShapeVec(tt); if (dstValid.size() != 2 || tmpValid.size() != 2) return emitOpError("expects tmp and dst to be rank-2 tiles"); @@ -11016,7 +11190,6 @@ mlir::LogicalResult mlir::pto::TPowOp::verify() { Type elem = getElemTy(baseTy); bool isIntElem = elem.isInteger(32) || elem.isInteger(16) || elem.isInteger(8); - bool isFpElem = elem.isF16() || elem.isF32() || elem.isBF16(); auto verifyA2A3 = [&]() -> LogicalResult { if (getPrecisionType() == pto::PowPrecision::HighPrecision) return emitOpError( @@ -11042,10 +11215,6 @@ mlir::LogicalResult mlir::pto::TPowOp::verify() { if (failed(dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5))) return failure(); - if (isFpElem && !getTmp()) - return emitOpError( - "expects tmp when element type is floating-point (required by the " - "floating-point pow lowering)"); if (isIntElem && getTmp()) return emitOpError( "does not accept tmp when element type is integer (the integer pow " @@ -11079,7 +11248,6 @@ mlir::LogicalResult mlir::pto::TPowSOp::verify() { // Same dtype matrix as TPowOp; see comment in TPowOp::verify. bool isIntElem = elem.isInteger(32) || elem.isInteger(16) || elem.isInteger(8); - bool isFpElem = elem.isF16() || elem.isF32() || elem.isBF16(); auto verifyA2A3 = [&]() -> LogicalResult { if (getPrecisionType() == pto::PowPrecision::HighPrecision) return emitOpError( @@ -11105,10 +11273,6 @@ mlir::LogicalResult mlir::pto::TPowSOp::verify() { if (failed(dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5))) return failure(); - if (isFpElem && !getTmp()) - return emitOpError( - "expects tmp when element type is floating-point (required by the " - "floating-point pow lowering)"); if (isIntElem && getTmp()) return emitOpError( "does not accept tmp when element type is integer (the integer pows " @@ -12235,8 +12399,267 @@ mlir::LogicalResult mlir::pto::TRowExpandMinOp::verify() { } +static ParseResult parseOptionalTmpRowReductionOp(OpAsmParser &parser, + OperationState &result) { + OpAsmParser::UnresolvedOperand src, tmp, dst; + Type srcTy, tmpTy, dstTy; + bool hasTmp = false; + + if (parser.parseKeyword("ins") || parser.parseLParen() || + parser.parseOperand(src)) + return failure(); + if (succeeded(parser.parseOptionalComma())) { + if (parser.parseOperand(tmp)) + return failure(); + hasTmp = true; + } + if (parser.parseColonType(srcTy)) + return failure(); + if (hasTmp && (parser.parseComma() || parser.parseType(tmpTy))) + return failure(); + if (parser.parseRParen() || parser.parseKeyword("outs") || + parser.parseLParen() || parser.parseOperand(dst) || + parser.parseColonType(dstTy) || parser.parseRParen() || + parser.parseOptionalAttrDict(result.attributes)) + return failure(); + + if (parser.resolveOperand(src, srcTy, result.operands)) + return failure(); + if (hasTmp && parser.resolveOperand(tmp, tmpTy, result.operands)) + return failure(); + if (parser.resolveOperand(dst, dstTy, result.operands)) + return failure(); + result.addAttribute( + "operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr({1, hasTmp ? 1 : 0, 1})); + return success(); +} + +static void printOptionalTmpRowReductionOp(OpAsmPrinter &p, Operation *op, + Value src, Value tmp, Value dst) { + p << " ins(" << src; + if (tmp) + p << ", " << tmp; + p << " : " << src.getType(); + if (tmp) + p << ", " << tmp.getType(); + p << ") outs(" << dst << " : " << dst.getType() << ")"; + p.printOptionalAttrDict(op->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); +} + +static ParseResult parseOptionalTmpFixedDpsOp( + OpAsmParser &parser, OperationState &result, unsigned minInputs, + unsigned maxInputs, ArrayRef noTmpSegments, + ArrayRef withTmpSegments) { + SmallVector inputs; + SmallVector inputTypes; + OpAsmParser::UnresolvedOperand dst; + Type dstType; + if (parser.parseKeyword("ins") || parser.parseLParen()) + return failure(); + do { + inputs.emplace_back(); + if (parser.parseOperand(inputs.back())) + return failure(); + } while (succeeded(parser.parseOptionalComma())); + if (inputs.size() < minInputs || inputs.size() > maxInputs || + parser.parseColon()) + return failure(); + for (unsigned i = 0; i < inputs.size(); ++i) { + if (i && parser.parseComma()) + return failure(); + Type type; + if (parser.parseType(type)) + return failure(); + inputTypes.push_back(type); + } + if (parser.parseRParen() || parser.parseKeyword("outs") || + parser.parseLParen() || parser.parseOperand(dst) || + parser.parseColonType(dstType) || parser.parseRParen() || + parser.parseOptionalAttrDict(result.attributes) || + parser.resolveOperands(inputs, inputTypes, parser.getCurrentLocation(), + result.operands) || + parser.resolveOperand(dst, dstType, result.operands)) + return failure(); + result.addAttribute( + "operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr( + inputs.size() == minInputs ? noTmpSegments : withTmpSegments)); + return success(); +} + +static void printOptionalTmpFixedDpsOp(OpAsmPrinter &p, Operation *op, + ArrayRef inputs, Value dst) { + p << " ins("; + llvm::interleaveComma(inputs, p, [&](Value value) { p << value; }); + p << " : "; + llvm::interleaveComma(inputs, p, + [&](Value value) { p << value.getType(); }); + p << ") outs(" << dst << " : " << dst.getType() << ")"; + p.printOptionalAttrDict(op->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); +} + +ParseResult mlir::pto::TTransOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpFixedDpsOp(parser, result, 1, 2, {1, 0, 1}, + {1, 1, 1}); +} +void mlir::pto::TTransOp::print(OpAsmPrinter &p) { + SmallVector inputs{getSrc()}; + if (getTmp()) + inputs.push_back(getTmp()); + printOptionalTmpFixedDpsOp(p, getOperation(), inputs, getDst()); +} + +ParseResult mlir::pto::TPReluOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpFixedDpsOp(parser, result, 2, 3, {1, 1, 0, 1}, + {1, 1, 1, 1}); +} +void mlir::pto::TPReluOp::print(OpAsmPrinter &p) { + SmallVector inputs{getSrc0(), getSrc1()}; + if (getTmp()) + inputs.push_back(getTmp()); + printOptionalTmpFixedDpsOp(p, getOperation(), inputs, getDst()); +} + +ParseResult mlir::pto::TRemOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpFixedDpsOp(parser, result, 2, 3, {1, 1, 0, 1}, + {1, 1, 1, 1}); +} +void mlir::pto::TRemOp::print(OpAsmPrinter &p) { + SmallVector inputs{getSrc0(), getSrc1()}; + if (getTmp()) + inputs.push_back(getTmp()); + printOptionalTmpFixedDpsOp(p, getOperation(), inputs, getDst()); +} + +ParseResult mlir::pto::TRemSOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpFixedDpsOp(parser, result, 2, 3, {1, 1, 0, 1}, + {1, 1, 1, 1}); +} +void mlir::pto::TRemSOp::print(OpAsmPrinter &p) { + SmallVector inputs{getSrc(), getScalar()}; + if (getTmp()) + inputs.push_back(getTmp()); + printOptionalTmpFixedDpsOp(p, getOperation(), inputs, getDst()); +} + +ParseResult mlir::pto::TSelOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpFixedDpsOp(parser, result, 3, 4, + {1, 1, 1, 0, 1}, {1, 1, 1, 1, 1}); +} +void mlir::pto::TSelOp::print(OpAsmPrinter &p) { + SmallVector inputs{getMask(), getSrc0(), getSrc1()}; + if (getTmp()) + inputs.push_back(getTmp()); + printOptionalTmpFixedDpsOp(p, getOperation(), inputs, getDst()); +} + +ParseResult mlir::pto::TSelSOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpFixedDpsOp(parser, result, 3, 4, + {1, 1, 0, 1, 1}, {1, 1, 1, 1, 1}); +} +void mlir::pto::TSelSOp::print(OpAsmPrinter &p) { + SmallVector inputs{getMask(), getSrc()}; + if (getTmp()) + inputs.push_back(getTmp()); + inputs.push_back(getScalar()); + printOptionalTmpFixedDpsOp(p, getOperation(), inputs, getDst()); +} + +ParseResult mlir::pto::TColArgMaxOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TColArgMaxOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TColArgMinOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TColArgMinOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TRowMaxOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TRowMaxOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TRowArgMaxOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TRowArgMaxOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TRowMinOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TRowMinOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TRowArgMinOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TRowArgMinOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TRowSumOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TRowSumOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TRowProdOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TRowProdOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + mlir::LogicalResult mlir::pto::TRowMaxOp::verify() { auto verifyByArch = [&]() -> LogicalResult { + if (!getTmp()) + return verifyTRowReductionNoTmpCommon( + *this, getSrc().getType(), getDst().getType(), + "expects element type to be i16/i32/f16/f32"); return verifyTRowReductionWithTmpCommon( *this, getSrc().getType(), getTmp().getType(), getDst().getType(), "expects element type to be i16/i32/f16/f32"); @@ -12245,6 +12668,9 @@ mlir::LogicalResult mlir::pto::TRowMaxOp::verify() { } mlir::LogicalResult mlir::pto::TRowArgMaxOp::verify() { + if (!getTmp()) + return verifyTRowArgReductionNoTmp(getOperation(), getSrc().getType(), + getDst().getType()); auto verifyA2A3 = [&]() -> LogicalResult { return verifyTRowArgReductionOpA2A3(*this, getSrc().getType(), getTmp().getType(), getDst().getType()); @@ -12260,6 +12686,10 @@ mlir::LogicalResult mlir::pto::TRowArgMaxOp::verify() { mlir::LogicalResult mlir::pto::TRowMinOp::verify() { auto verifyByArch = [&]() -> LogicalResult { + if (!getTmp()) + return verifyTRowReductionNoTmpCommon( + *this, getSrc().getType(), getDst().getType(), + "expects element type to be i16/i32/f16/f32"); return verifyTRowReductionWithTmpCommon( *this, getSrc().getType(), getTmp().getType(), getDst().getType(), "expects element type to be i16/i32/f16/f32"); @@ -12268,6 +12698,9 @@ mlir::LogicalResult mlir::pto::TRowMinOp::verify() { } mlir::LogicalResult mlir::pto::TRowArgMinOp::verify() { + if (!getTmp()) + return verifyTRowArgReductionNoTmp(getOperation(), getSrc().getType(), + getDst().getType()); auto verifyA2A3 = [&]() -> LogicalResult { return verifyTRowArgReductionOpA2A3(*this, getSrc().getType(), getTmp().getType(), getDst().getType()); @@ -12283,6 +12716,10 @@ mlir::LogicalResult mlir::pto::TRowArgMinOp::verify() { mlir::LogicalResult mlir::pto::TRowSumOp::verify() { auto verifyByArch = [&]() -> LogicalResult { + if (!getTmp()) + return verifyTRowReductionNoTmpCommon( + *this, getSrc().getType(), getDst().getType(), + "expects element type to be i16/i32/f16/f32"); return verifyTRowReductionWithTmpCommon( *this, getSrc().getType(), getTmp().getType(), getDst().getType(), "expects element type to be i16/i32/f16/f32"); @@ -12292,11 +12729,19 @@ mlir::LogicalResult mlir::pto::TRowSumOp::verify() { mlir::LogicalResult mlir::pto::TRowProdOp::verify() { auto verifyA2A3 = [&]() -> LogicalResult { + if (!getTmp()) + return verifyTRowReductionNoTmpCommon( + *this, getSrc().getType(), getDst().getType(), + "expects A2/A3 trowprod element type to be i16/i32/f16/f32"); return verifyTRowReductionWithTmpCommon( *this, getSrc().getType(), getTmp().getType(), getDst().getType(), "expects A2/A3 trowprod element type to be i16/i32/f16/f32"); }; auto verifyA5 = [&]() -> LogicalResult { + if (!getTmp()) + return verifyTRowReductionNoTmpCommon( + *this, getSrc().getType(), getDst().getType(), + "expects A5 trowprod element type to be i16/i32/f16/f32"); return verifyTRowReductionWithTmpCommon( *this, getSrc().getType(), getTmp().getType(), getDst().getType(), "expects A5 trowprod element type to be i16/i32/f16/f32"); @@ -12555,16 +13000,17 @@ mlir::LogicalResult mlir::pto::TSelSOp::verify() { auto verifyCommon = [&]() -> FailureOr { Type tMask = getMask().getType(); Type tSrc = getSrc().getType(); - Type tTmp = getTmp().getType(); + Type tTmp = getTmp() ? getTmp().getType() : Type{}; Type tDst = getDst().getType(); if (failed(verifyTileBufCommon(*this, tMask, "mask")) || failed(verifyTileBufCommon(*this, tSrc, "src")) || - failed(verifyTileBufCommon(*this, tTmp, "tmp")) || failed(verifyTileBufCommon(*this, tDst, "dst"))) return failure(); + if (tTmp && failed(verifyTileBufCommon(*this, tTmp, "tmp"))) + return failure(); Type eMask = getElemTy(tMask), eSrc = getElemTy(tSrc); - Type eTmp = getElemTy(tTmp), eDst = getElemTy(tDst); - if (!eMask || !eSrc || !eTmp || !eDst) { + Type eDst = getElemTy(tDst); + if (!eMask || !eSrc || !eDst) { emitOpError("failed to get element type for operands"); return failure(); } @@ -12839,17 +13285,34 @@ mlir::LogicalResult mlir::pto::TSubSCOp::verify() { return emitOpError() << "expects src0, src1, and dst to have the same rank"; return mlir::success(); } +static bool ttransUsesTmp(Type srcTy, Type dstTy) { + auto srcShape = getShapeVec(srcTy); + auto dstShape = getShapeVec(dstTy); + unsigned elemBytes = getPTOStorageElemByteSize(getElemTy(srcTy)); + if (srcShape.size() != 2 || dstShape.size() != 2 || elemBytes == 0 || + llvm::is_contained(srcShape, ShapedType::kDynamic) || + llvm::is_contained(dstShape, ShapedType::kDynamic)) + return true; + int64_t rowStride = elemBytes == 1 ? 32 : 16; + int64_t elemPerBlock = 32 / elemBytes; + int64_t srcStride = srcShape[1]; + int64_t dstStride = dstShape[1]; + return dstStride % rowStride == 0 && srcStride % elemPerBlock == 0 && + srcStride / elemPerBlock <= 255; +} + mlir::LogicalResult mlir::pto::TTransOp::verify() { auto verifyA2A3 = [&]() -> LogicalResult { Type srcTy = getSrc().getType(); - Type tmpTy = getTmp().getType(); + Type tmpTy = getTmp() ? getTmp().getType() : Type{}; Type dstTy = getDst().getType(); if (failed(verifyTileBufCommon(*this, srcTy, "src")) || - failed(verifyTileBufCommon(*this, tmpTy, "tmp")) || failed(verifyTileBufCommon(*this, dstTy, "dst"))) return failure(); + if (tmpTy && failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) + return failure(); Type srcElem = getElemTy(srcTy); - Type tmpElem = getElemTy(tmpTy); + Type tmpElem = tmpTy ? getElemTy(tmpTy) : srcElem; Type dstElem = getElemTy(dstTy); if (!srcElem || !tmpElem || !dstElem || srcElem != dstElem || srcElem != tmpElem) return emitOpError() << "expects src and dst to have the same element type"; @@ -12875,14 +13338,15 @@ mlir::LogicalResult mlir::pto::TTransOp::verify() { }; auto verifyA5 = [&]() -> LogicalResult { Type srcTy = getSrc().getType(); - Type tmpTy = getTmp().getType(); + Type tmpTy = getTmp() ? getTmp().getType() : Type{}; Type dstTy = getDst().getType(); if (failed(verifyTileBufCommon(*this, srcTy, "src")) || - failed(verifyTileBufCommon(*this, tmpTy, "tmp")) || failed(verifyTileBufCommon(*this, dstTy, "dst"))) return failure(); + if (tmpTy && failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) + return failure(); Type srcElem = getElemTy(srcTy); - Type tmpElem = getElemTy(tmpTy); + Type tmpElem = tmpTy ? getElemTy(tmpTy) : srcElem; Type dstElem = getElemTy(dstTy); if (!srcElem || !tmpElem || !dstElem || srcElem != dstElem || srcElem != tmpElem) return emitOpError() << "expects src, tmp, and dst to have the same element type"; @@ -12920,6 +13384,100 @@ mlir::LogicalResult mlir::pto::TTransOp::verify() { return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); } +ParseResult mlir::pto::TXorOp::parse(OpAsmParser &parser, + OperationState &result) { + OpAsmParser::UnresolvedOperand src0, src1, tmp, dst; + Type src0Ty, src1Ty, tmpTy, dstTy; + bool hasTmp = false; + if (parser.parseKeyword("ins") || parser.parseLParen() || + parser.parseOperand(src0) || parser.parseComma() || + parser.parseOperand(src1)) + return failure(); + if (succeeded(parser.parseOptionalComma())) { + if (parser.parseOperand(tmp)) + return failure(); + hasTmp = true; + } + if (parser.parseColonType(src0Ty) || parser.parseComma() || + parser.parseType(src1Ty)) + return failure(); + if (hasTmp && (parser.parseComma() || parser.parseType(tmpTy))) + return failure(); + if (parser.parseRParen() || parser.parseKeyword("outs") || + parser.parseLParen() || parser.parseOperand(dst) || + parser.parseColonType(dstTy) || parser.parseRParen() || + parser.parseOptionalAttrDict(result.attributes)) + return failure(); + if (parser.resolveOperand(src0, src0Ty, result.operands) || + parser.resolveOperand(src1, src1Ty, result.operands) || + (hasTmp && parser.resolveOperand(tmp, tmpTy, result.operands)) || + parser.resolveOperand(dst, dstTy, result.operands)) + return failure(); + result.addAttribute( + "operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr({1, 1, hasTmp ? 1 : 0, 1})); + return success(); +} + +void mlir::pto::TXorOp::print(OpAsmPrinter &p) { + p << " ins(" << getSrc0() << ", " << getSrc1(); + if (getTmp()) + p << ", " << getTmp(); + p << " : " << getSrc0().getType() << ", " << getSrc1().getType(); + if (getTmp()) + p << ", " << getTmp().getType(); + p << ") outs(" << getDst() << " : " << getDst().getType() << ")"; + p.printOptionalAttrDict((*this)->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); +} + +ParseResult mlir::pto::TXorSOp::parse(OpAsmParser &parser, + OperationState &result) { + OpAsmParser::UnresolvedOperand src, scalar, tmp, dst; + Type srcTy, scalarTy, tmpTy, dstTy; + bool hasTmp = false; + if (parser.parseKeyword("ins") || parser.parseLParen() || + parser.parseOperand(src) || parser.parseComma() || + parser.parseOperand(scalar)) + return failure(); + if (succeeded(parser.parseOptionalComma())) { + if (parser.parseOperand(tmp)) + return failure(); + hasTmp = true; + } + if (parser.parseColonType(srcTy) || parser.parseComma() || + parser.parseType(scalarTy)) + return failure(); + if (hasTmp && (parser.parseComma() || parser.parseType(tmpTy))) + return failure(); + if (parser.parseRParen() || parser.parseKeyword("outs") || + parser.parseLParen() || parser.parseOperand(dst) || + parser.parseColonType(dstTy) || parser.parseRParen() || + parser.parseOptionalAttrDict(result.attributes)) + return failure(); + if (parser.resolveOperand(src, srcTy, result.operands) || + parser.resolveOperand(scalar, scalarTy, result.operands) || + (hasTmp && parser.resolveOperand(tmp, tmpTy, result.operands)) || + parser.resolveOperand(dst, dstTy, result.operands)) + return failure(); + result.addAttribute( + "operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr({1, 1, hasTmp ? 1 : 0, 1})); + return success(); +} + +void mlir::pto::TXorSOp::print(OpAsmPrinter &p) { + p << " ins(" << getSrc() << ", " << getScalar(); + if (getTmp()) + p << ", " << getTmp(); + p << " : " << getSrc().getType() << ", " << getScalar().getType(); + if (getTmp()) + p << ", " << getTmp().getType(); + p << ") outs(" << getDst() << " : " << getDst().getType() << ")"; + p.printOptionalAttrDict((*this)->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); +} + mlir::LogicalResult mlir::pto::TXorOp::verify() { auto verifyBase = [&]() -> FailureOr { return verifyMatchingRowMajorBinaryTileOpCommon( @@ -12931,16 +13489,20 @@ mlir::LogicalResult mlir::pto::TXorOp::verify() { FailureOr elemOr = verifyBase(); if (failed(elemOr)) return failure(); - Type tmpTy = getTmp().getType(); - if (failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) - return failure(); Type elem = *elemOr; - if (getElemTy(tmpTy) != elem) - return emitOpError("expects tmp to have the same element type as src0, src1, and dst"); - if (!isRowMajorTileBuf(tmpTy)) - return emitOpError("expects tmp to use row-major layout"); - if (failed(verifyTileBufSameValidShape(*this, tmpTy, getDst().getType(), "tmp", "dst"))) - return failure(); + if (getTmp()) { + Type tmpTy = getTmp().getType(); + if (failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) + return failure(); + if (getElemTy(tmpTy) != elem) + return emitOpError( + "expects tmp to have the same element type as src0, src1, and dst"); + if (!isRowMajorTileBuf(tmpTy)) + return emitOpError("expects tmp to use row-major layout"); + if (failed(verifyTileBufSameValidShape( + *this, tmpTy, getDst().getType(), "tmp", "dst"))) + return failure(); + } auto it = mlir::dyn_cast(elem); if (!it || (it.getWidth() != 8 && it.getWidth() != 16 && it.getWidth() != 32)) @@ -12975,14 +13537,17 @@ mlir::LogicalResult mlir::pto::TXorSOp::verify() { FailureOr elemOr = verifyCommon(); if (failed(elemOr)) return failure(); - Type tmpTy = getTmp().getType(); - if (failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) - return failure(); Type elem = *elemOr; - if (getElemTy(tmpTy) != elem) - return emitOpError("expects tmp to have the same element type as src and dst"); - if (!isRowMajorTileBuf(tmpTy)) - return emitOpError("expects tmp to use row-major layout"); + if (getTmp()) { + Type tmpTy = getTmp().getType(); + if (failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) + return failure(); + if (getElemTy(tmpTy) != elem) + return emitOpError( + "expects tmp to have the same element type as src and dst"); + if (!isRowMajorTileBuf(tmpTy)) + return emitOpError("expects tmp to use row-major layout"); + } auto it = mlir::dyn_cast(elem); if (!it || (it.getWidth() != 8 && it.getWidth() != 16)) return emitOpError( @@ -14388,16 +14953,22 @@ PTO_DEFINE_UNARY_EFFECTS(TColProdOp, getSrcMutable(), getDstMutable()) void TColArgMaxOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TColArgMinOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14406,6 +14977,7 @@ void TColSumOp::getEffects( PTO_ADD_READ(getSrcMutable()); auto tmp = getTmpMutable(); if (!tmp.empty()) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); } PTO_ADD_WRITE(getDstMutable()); @@ -14414,6 +14986,11 @@ void TColSumOp::getEffects( void TCvtOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TRandomOp::getEffects( @@ -14482,8 +15059,11 @@ void TGatherOp::getEffects( PTO_ADD_WRITE(cdst[0]); if (auto indices = getIndicesMutable(); !indices.empty()) PTO_ADD_READ(indices[0]); - if (auto tmp = getTmpMutable(); !tmp.empty()) + if (auto tmp = getTmpMutable(); + !tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14504,8 +15084,10 @@ void TMrgSortOp::getEffects( PTO_ADD_READ(opnd); } auto tmp = getTmpMutable(); - if (!tmp.empty()) + if (!tmp.empty()) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } for (auto &opnd : getDstsMutable()) { PTO_ADD_WRITE(opnd); } @@ -14552,8 +15134,11 @@ void TPReluOp::getEffects( // A5 pto-isa TPRELU implementation does not consume tmp; modeling tmp as a // write-only scratch on A5 incorrectly inflates local-memory planning and // can trigger false vec-overflow diagnostics. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14565,8 +15150,10 @@ void TQuantOp::getEffects( if (!offsetRange.empty()) PTO_ADD_READ(offsetRange[0]); auto tmpRange = getTmpMutable(); - if (!tmpRange.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmpRange.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmpRange[0]); PTO_ADD_WRITE(tmpRange[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14591,16 +15178,22 @@ void TRemOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TRemSOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14609,8 +15202,10 @@ void TPowOp::getEffects( PTO_ADD_READ(getBaseMutable()); PTO_ADD_READ(getExpMutable()); auto tmp = getTmpMutable(); - if (!tmp.empty()) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14618,8 +15213,10 @@ void TPowSOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); auto tmp = getTmpMutable(); - if (!tmp.empty()) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } PTO_DEFINE_UNARY_EFFECTS(TRowExpandOp, getSrcMutable(), getDstMutable()) @@ -14710,8 +15307,11 @@ void TRowExpandMinOp::getEffects( void TRowMaxOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14721,16 +15321,22 @@ void TRowArgMaxOp::getEffects( // A5 lowering does not consume tmp for TROWARGMAX; modeling tmp as a // scratch write inflates local-memory planning and can trigger false // vec-overflow diagnostics, mirroring the fixed A5 TPRELU issue. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TRowMinOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14740,32 +15346,43 @@ void TRowArgMinOp::getEffects( // A5 lowering does not consume tmp for TROWARGMIN; modeling tmp as a // scratch write inflates local-memory planning and can trigger false // vec-overflow diagnostics, mirroring the fixed A5 TPRELU issue. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TRowSumOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TRowProdOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TRsqrtOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); auto tmp = getTmpMutable(); - if (!tmp.empty()) + if (!tmp.empty()) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14789,8 +15406,11 @@ void TSelOp::getEffects( // A5 lowering does not consume tmp for TSEL; modeling tmp as a scratch // write inflates local-memory planning and can trigger false vec-overflow // diagnostics. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14802,8 +15422,11 @@ void TSelSOp::getEffects( // A5 lowering does not consume tmp for TSELS; modeling tmp as a scratch // write inflates local-memory planning and can trigger false vec-overflow // diagnostics. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14836,8 +15459,11 @@ void TXorSOp::getEffects( // A5 lowering does not consume tmp for TXORS; modeling tmp as a scratch // write inflates local-memory planning and can trigger false vec-overflow // diagnostics. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14849,8 +15475,11 @@ void TXorOp::getEffects( // A5 lowering does not consume tmp for TXOR; modeling tmp as a scratch // write inflates local-memory planning and can trigger false vec-overflow // diagnostics. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14858,7 +15487,11 @@ void TXorOp::getEffects( void TTransOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && ttransUsesTmp(getSrc().getType(), getDst().getType())) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } diff --git a/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp b/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp index 1709f10e4b..6ee5b20d26 100644 --- a/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp +++ b/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp @@ -19,6 +19,7 @@ #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Matchers.h" +#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "llvm/Support/Casting.h" #include "llvm/Support/ErrorHandling.h" @@ -40,6 +41,44 @@ namespace { static constexpr uint64_t kVectorRegisterSizeInBytes = 256U; static constexpr unsigned kPipeVPruneMinRepeat = 16U; +static bool hasReadWriteScratchDependency( + Operation *op, const DepBaseMemInfoPairVec &dependencies) { + auto effectsOp = dyn_cast_or_null(op); + if (!effectsOp) + return false; + + llvm::DenseSet reads; + llvm::DenseSet writes; + SmallVector, 8> effects; + effectsOp.getEffects(effects); + for (const auto &effect : effects) { + Value value = effect.getValue(); + if (!value) + continue; + if (isa(effect.getEffect())) + reads.insert(value); + if (isa(effect.getEffect())) + writes.insert(value); + } + + ValueRange dpsInits; + if (auto ptoDpsOp = dyn_cast(op)) + dpsInits = ptoDpsOp.getDpsInits(); + else if (auto dpsOp = dyn_cast(op)) + dpsInits = dpsOp.getDpsInits(); + return llvm::any_of(writes, [&](Value value) { + if (!reads.contains(value) || llvm::is_contained(dpsInits, value)) + return false; + return llvm::any_of(dependencies, [&](const auto &dependency) { + auto matches = [&](const BaseMemInfo *info) { + return info && + (info->baseBuffer == value || info->rootBuffer == value); + }; + return matches(dependency.first) || matches(dependency.second); + }); + }); +} + struct RepeatAccessShape { SmallVector fullShape; SmallVector validShape; @@ -511,6 +550,16 @@ bool InsertSyncAnalysis::CanPrunePipeVBarrier( return false; } + // The same-access fast path only applies to a producer output consumed by + // the next op. A read/write non-DPS operand is scratch state; pruning its + // WAW dependency would allow two vector instructions to use it concurrently. + if (hasReadWriteScratchDependency(nowCompound->elementOp, + depBaseMemInfosVec) || + hasReadWriteScratchDependency(frontCompound->elementOp, + depBaseMemInfosVec)) { + return false; + } + // PIPE_V has a hardware-safe same-access chain case: exact same-access // dependencies from the producer result to the consumer source do not require // a vector-pipe barrier once the producer repeat is large enough. Keep the diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp index abb14a89e2..8c0d60eef7 100644 --- a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -67,6 +67,64 @@ static SmallVector getValidShapeVec(Type ty) { return {}; } +static SmallVector getShapeVec(Type ty) { + if (auto tileTy = dyn_cast(ty)) + return SmallVector(tileTy.getShape().begin(), + tileTy.getShape().end()); + return {}; +} + +static int64_t ceilDiv(int64_t lhs, int64_t rhs) { + return (lhs + rhs - 1) / rhs; +} + +static bool hasDynamicDim(ArrayRef dims) { + return llvm::any_of(dims, [](int64_t dim) { + return dim == ShapedType::kDynamic; + }); +} + +static pto::TileBufType makeVecTmpType(MLIRContext *ctx, + ArrayRef shape, + Type elementType, + ArrayRef validShape) { + return pto::TileBufType::get( + ctx, shape, elementType, + pto::AddressSpaceAttr::get(ctx, pto::AddressSpace::VEC), validShape, + makeRowMajorNoneBoxConfig(ctx)); +} + +static FailureOr makeSameShapeTmpType(MLIRContext *ctx, + Value like, + Type elementType = {}) { + auto likeTy = dyn_cast(like.getType()); + if (!likeTy) + return failure(); + if (!elementType) + elementType = likeTy.getElementType(); + auto shape = getShapeVec(like.getType()); + auto validShape = getValidShapeVec(like.getType()); + if (shape.empty() || validShape.empty() || hasDynamicDim(shape) || + hasDynamicDim(validShape)) + return failure(); + return makeVecTmpType(ctx, shape, elementType, validShape); +} + +static FailureOr createAllocTmp(OpBuilder &builder, Location loc, + pto::TileBufType tmpType) { + return builder + .create(loc, tmpType, Value(), Value(), Value()) + .getResult(); +} + +static void copyAttrsExceptOperandSegments(Operation *from, OperationState &to) { + for (NamedAttribute attr : from->getAttrs()) { + if (attr.getName() == "operandSegmentSizes") + continue; + to.addAttribute(attr.getName(), attr.getValue()); + } +} + static bool validShapesCompatible(ArrayRef lhs, ArrayRef rhs) { if (lhs.size() != rhs.size()) @@ -197,6 +255,553 @@ static LogicalResult materializeTRowExpandTmp(OpTy op, bool requireExplicitTmp, return success(); } +static LogicalResult replaceTColSumWithTmp(pto::TColSumOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp() || !op.getIsBinary()) + return success(); + if (requireExplicitTmp) + return op.emitOpError( + "requires explicit tmp for binary tcolsum when PlanMemory is skipped"); + + auto srcTy = dyn_cast(op.getSrc().getType()); + if (!srcTy) + return op.emitOpError("expects tile_buf src when materializing implicit tmp"); + auto valid = getValidShapeVec(op.getSrc().getType()); + if (valid.size() != 2 || hasDynamicDim(valid)) + return op.emitOpError( + "requires static src valid_shape to materialize binary tcolsum tmp"); + + SmallVector tmpShape{ceilDiv(valid[0], 2), valid[1]}; + auto tmpType = makeVecTmpType(ctx, tmpShape, srcTy.getElementType(), tmpShape); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getSrc(), *tmp, op.getDst()}); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceTQuantWithTmp(pto::TQuantOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp() || pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + + FailureOr tmpType = makeSameShapeTmpType( + ctx, op.getSrc(), Float32Type::get(ctx)); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf src to materialize implicit tquant tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + SmallVector operands{op.getSrc(), op.getFp()}; + if (op.getOffset()) + operands.push_back(op.getOffset()); + operands.push_back(*tmp); + operands.push_back(op.getDst()); + state.addOperands(operands); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr( + {1, 1, op.getOffset() ? 1 : 0, 1, 1})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static bool isFloatingPointTile(Value value) { + auto tileTy = dyn_cast(value.getType()); + return tileTy && isa(tileTy.getElementType()); +} + +static LogicalResult replaceTPowWithTmp(pto::TPowOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp() || !isFloatingPointTile(op.getDst())) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + + FailureOr tmpType = makeSameShapeTmpType(ctx, op.getDst()); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf dst to materialize implicit tpow tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getBase(), op.getExp(), op.getDst(), *tmp}); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceTPowSWithTmp(pto::TPowSOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp() || !isFloatingPointTile(op.getDst())) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + + FailureOr tmpType = makeSameShapeTmpType(ctx, op.getDst()); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf dst to materialize implicit tpows tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getSrc(), op.getScalar(), op.getDst(), *tmp}); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceTGatherWithTmp(pto::TGatherOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp() || op.hasMaskForm()) + return success(); + if (!op.hasIndexForm() && !op.hasCompareForm()) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + + FailureOr tmpType = failure(); + if (op.hasIndexForm()) { + tmpType = makeSameShapeTmpType(ctx, op.getIndices()); + } else { + auto srcTy = dyn_cast(op.getSrc().getType()); + auto dstTy = dyn_cast(op.getDst().getType()); + if (!srcTy || !dstTy) + return op.emitOpError( + "expects tile_buf operands when materializing compare-form tgather tmp"); + auto srcShape = getShapeVec(op.getSrc().getType()); + if (srcShape.size() != 2 || hasDynamicDim(srcShape)) + return op.emitOpError( + "requires static src shape to materialize compare-form tgather tmp"); + int64_t bytes = srcShape[0] * srcShape[1] * 4 + srcShape[0] * 4; + tmpType = makeVecTmpType(ctx, {1, bytes}, IntegerType::get(ctx, 8), + {1, bytes}); + } + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf indices/src to materialize implicit tgather tmp"); + + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + SmallVector operands{op.getSrc(), op.getDst()}; + if (op.getCdst()) + operands.push_back(op.getCdst()); + if (op.getIndices()) + operands.push_back(op.getIndices()); + operands.push_back(*tmp); + if (op.getKValue()) + operands.push_back(op.getKValue()); + state.addOperands(operands); + state.addAttribute( + "operandSegmentSizes", + builder.getDenseI32ArrayAttr({1, 1, op.getCdst() ? 1 : 0, + op.getIndices() ? 1 : 0, 1, + op.getKValue() ? 1 : 0})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceTSort32WithTmp(pto::TSort32Op op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp()) + return success(); + auto valid = getValidShapeVec(op.getSrc().getType()); + if (valid.size() != 2 || valid[1] == ShapedType::kDynamic || + valid[1] % 32 == 0) + return success(); + if (requireExplicitTmp) + return op.emitOpError( + "requires explicit tmp for non-32-aligned tsort32 when PlanMemory is skipped"); + + FailureOr tmpType = makeSameShapeTmpType(ctx, op.getSrc()); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf src to materialize implicit tsort32 tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getSrc(), op.getIdx(), *tmp, op.getDst()}); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr({1, 1, 1, 1})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +template +static LogicalResult replaceRowReductionWithTmp(OpTy op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp()) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + + FailureOr tmpType = makeSameShapeTmpType(ctx, op.getSrc()); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf src to materialize implicit row-reduction tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getSrc(), *tmp, op.getDst()}); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr({1, 1, 1})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceTXorWithTmp(pto::TXorOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp()) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + FailureOr tmpType = makeSameShapeTmpType(ctx, op.getDst()); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf dst to materialize implicit txor tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getSrc0(), op.getSrc1(), *tmp, op.getDst()}); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr({1, 1, 1, 1})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceTXorSWithTmp(pto::TXorSOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp()) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + FailureOr tmpType = makeSameShapeTmpType(ctx, op.getDst()); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf dst to materialize implicit txors tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getSrc(), op.getScalar(), *tmp, op.getDst()}); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr({1, 1, 1, 1})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceFixedDpsOpWithTmp( + Operation *op, ArrayRef operands, pto::TileBufType tmpType, + ArrayRef operandSegments, bool requireExplicitTmp, + StringRef opName) { + if (requireExplicitTmp) + return op->emitOpError( + "requires explicit tmp when PlanMemory is skipped"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op->getLoc(), tmpType); + if (failed(tmp)) + return failure(); + SmallVector finalOperands; + finalOperands.reserve(operands.size() + 1); + for (Value operand : operands) { + if (operand) + finalOperands.push_back(operand); + else + finalOperands.push_back(*tmp); + } + OperationState state(op->getLoc(), op->getName()); + state.addOperands(finalOperands); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr(operandSegments)); + copyAttrsExceptOperandSegments(op, state); + builder.create(state); + op->erase(); + (void)opName; + return success(); +} + +static FailureOr makeTPReluTmpType(MLIRContext *ctx, + Value dst) { + auto dstTy = dyn_cast(dst.getType()); + auto shape = getShapeVec(dst.getType()); + auto valid = getValidShapeVec(dst.getType()); + if (!dstTy || shape.size() != 2 || valid.size() != 2 || + hasDynamicDim(shape) || hasDynamicDim(valid)) + return failure(); + int64_t validCols = ceilDiv(valid[1], 8); + int64_t cols = std::max(32, ceilDiv(validCols, 32) * 32); + return makeVecTmpType(ctx, {valid[0] + 1, cols}, IntegerType::get(ctx, 8), + {valid[0], validCols}); +} + +static FailureOr makeRowsTmpType(MLIRContext *ctx, + Value dst, int64_t rows) { + auto dstTy = dyn_cast(dst.getType()); + auto shape = getShapeVec(dst.getType()); + auto valid = getValidShapeVec(dst.getType()); + if (!dstTy || shape.size() != 2 || valid.size() != 2 || + hasDynamicDim(shape) || hasDynamicDim(valid)) + return failure(); + return makeVecTmpType(ctx, {rows, shape[1]}, dstTy.getElementType(), + {rows, valid[1]}); +} + +static LogicalResult materializeFixedMandatoryTmp(Operation *op, + bool requireExplicitTmp, + MLIRContext *ctx) { + return llvm::TypeSwitch(op) + .Case([&](auto typedOp) -> LogicalResult { + if (typedOp.getTmp()) + return success(); + auto type = makeTPReluTmpType(ctx, typedOp.getDst()); + if (failed(type)) + return typedOp.emitOpError( + "requires static tile_buf dst to materialize implicit tprelu tmp"); + return replaceFixedDpsOpWithTmp( + op, {typedOp.getSrc0(), typedOp.getSrc1(), Value(), + typedOp.getDst()}, + *type, {1, 1, 1, 1}, requireExplicitTmp, "tprelu"); + }) + .Case([&](auto typedOp) -> LogicalResult { + if (typedOp.getTmp()) + return success(); + auto type = makeRowsTmpType(ctx, typedOp.getDst(), 2); + if (failed(type)) + return typedOp.emitOpError( + "requires static tile_buf dst to materialize implicit trem tmp"); + return replaceFixedDpsOpWithTmp( + op, {typedOp.getSrc0(), typedOp.getSrc1(), Value(), + typedOp.getDst()}, + *type, {1, 1, 1, 1}, requireExplicitTmp, "trem"); + }) + .Case([&](auto typedOp) -> LogicalResult { + if (typedOp.getTmp()) + return success(); + auto type = makeRowsTmpType(ctx, typedOp.getDst(), 1); + if (failed(type)) + return typedOp.emitOpError( + "requires static tile_buf dst to materialize implicit trems tmp"); + return replaceFixedDpsOpWithTmp( + op, {typedOp.getSrc(), typedOp.getScalar(), Value(), + typedOp.getDst()}, + *type, {1, 1, 1, 1}, requireExplicitTmp, "trems"); + }) + .Case([&](auto typedOp) -> LogicalResult { + if (typedOp.getTmp()) + return success(); + auto type = makeVecTmpType(ctx, {1, 16}, IntegerType::get(ctx, 32), + {1, 16}); + return replaceFixedDpsOpWithTmp( + op, {typedOp.getMask(), typedOp.getSrc0(), typedOp.getSrc1(), + Value(), typedOp.getDst()}, + type, {1, 1, 1, 1, 1}, requireExplicitTmp, "tsel"); + }) + .Case([&](auto typedOp) -> LogicalResult { + if (typedOp.getTmp()) + return success(); + auto type = makeRowsTmpType(ctx, typedOp.getSrc(), 1); + if (failed(type)) + return typedOp.emitOpError( + "requires static tile_buf src to materialize implicit tsels tmp"); + return replaceFixedDpsOpWithTmp( + op, {typedOp.getMask(), typedOp.getSrc(), Value(), + typedOp.getScalar(), typedOp.getDst()}, + *type, {1, 1, 1, 1, 1}, requireExplicitTmp, "tsels"); + }) + .Case([&](auto typedOp) -> LogicalResult { + if (typedOp.getTmp()) + return success(); + auto srcTy = dyn_cast(typedOp.getSrc().getType()); + auto dstTy = dyn_cast(typedOp.getDst().getType()); + auto srcShape = getShapeVec(typedOp.getSrc().getType()); + auto dstShape = getShapeVec(typedOp.getDst().getType()); + if (!srcTy || !dstTy || srcShape.size() != 2 || dstShape.size() != 2 || + hasDynamicDim(srcShape) || hasDynamicDim(dstShape)) + return typedOp.emitOpError( + "requires static tile_buf src to materialize implicit ttrans tmp"); + auto elemBytes = getElemBytes(srcTy.getElementType()); + if (!elemBytes) + return typedOp.emitOpError("failed to infer ttrans element size"); + int64_t rowStride = *elemBytes == 1 ? 32 : 16; + int64_t elemPerBlock = 32 / *elemBytes; + bool usesTmp = dstShape[1] % rowStride == 0 && + srcShape[1] % elemPerBlock == 0 && + srcShape[1] / elemPerBlock <= 255; + FailureOr type = makeSameShapeTmpType( + ctx, typedOp.getSrc()); + if (!usesTmp) + type = makeVecTmpType(ctx, {1, elemPerBlock}, + srcTy.getElementType(), {1, elemPerBlock}); + if (failed(type)) + return typedOp.emitOpError("failed to build implicit ttrans tmp"); + return replaceFixedDpsOpWithTmp( + op, {typedOp.getSrc(), Value(), typedOp.getDst()}, *type, + {1, 1, 1}, requireExplicitTmp, "ttrans"); + }) + .Default([](Operation *) { return success(); }); +} + +static bool tcvtNeedsTmp(pto::TCvtOp op) { + if (pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5 || + op.getSatMode() != pto::SaturationMode::OFF) + return false; + auto srcTy = dyn_cast(op.getSrc().getType()); + auto dstTy = dyn_cast(op.getDst().getType()); + if (!srcTy || !dstTy) + return false; + Type srcElem = srcTy.getElementType(); + Type dstElem = dstTy.getElementType(); + return (srcElem.isF32() && dstElem.isInteger(16)) || + (srcElem.isF16() && + (dstElem.isInteger(16) || dstElem.isInteger(8))); +} + +static FailureOr makeTCvtTmpType(MLIRContext *ctx, + pto::TCvtOp op) { + auto srcShape = getShapeVec(op.getSrc().getType()); + auto dstValid = getValidShapeVec(op.getDst().getType()); + auto srcTy = dyn_cast(op.getSrc().getType()); + auto dstTy = dyn_cast(op.getDst().getType()); + if (!srcTy || !dstTy || srcShape.size() != 2 || dstValid.size() != 2 || + hasDynamicDim(srcShape) || hasDynamicDim(dstValid)) + return failure(); + int64_t rows = dstValid[0], cols = dstValid[1]; + int64_t bytes = 0; + if (rows > 0 && cols > 0 && srcTy.getElementType().isF32()) { + int64_t head = 4 * 64 * std::min(cols / 64, 255); + int64_t remainder = cols % 64; + int64_t tail = remainder == 0 + ? 0 + : 32 * ((std::min(rows, 255) - 1) * + (srcShape[1] / 8) + + ceilDiv(remainder, 8)); + bytes = std::max(head, tail); + } else if (cols > 0 && srcTy.getElementType().isF16()) { + int64_t width = std::min(cols, 64); + int64_t halfToI16 = 32 * ceilDiv(width, 8); + int64_t halfToI8 = std::max(halfToI16, 128 + 32 * ceilDiv(width, 16)); + bytes = dstTy.getElementType().isInteger(8) ? halfToI8 : halfToI16; + } + int64_t allocatedBytes = std::max(32, ceilDiv(bytes, 32) * 32); + return makeVecTmpType(ctx, {1, allocatedBytes}, IntegerType::get(ctx, 8), + {1, allocatedBytes}); +} + +static LogicalResult materializeTCvtTmp(pto::TCvtOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp() || !tcvtNeedsTmp(op)) + return success(); + if (requireExplicitTmp) + return op.emitOpError( + "requires explicit tmp for non-saturating narrowing tcvt when PlanMemory is skipped"); + auto type = makeTCvtTmpType(ctx, op); + if (failed(type)) + return op.emitOpError( + "requires static tile_buf shapes to materialize implicit tcvt tmp"); + return replaceFixedDpsOpWithTmp(op.getOperation(), + {op.getSrc(), Value(), op.getDst()}, *type, + {1, 1, 1}, requireExplicitTmp, "tcvt"); +} + +static LogicalResult materializeTMrgSortTmp(pto::TMrgSortOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (!op.isFormat2WithoutTmp()) + return success(); + if (requireExplicitTmp) + return op.emitOpError( + "requires explicit tmp for tmrgsort format2 when PlanMemory is skipped"); + int64_t totalCols = 0; + Type elementType; + SmallVector operands; + for (Value src : op.getSrcs()) { + auto srcTy = dyn_cast(src.getType()); + auto shape = getShapeVec(src.getType()); + if (!srcTy || shape.size() != 2 || hasDynamicDim(shape)) + return op.emitOpError( + "requires static rank-2 tile_buf srcs to materialize tmrgsort tmp"); + if (!elementType) + elementType = srcTy.getElementType(); + totalCols += shape[1]; + operands.push_back(src); + } + if (!elementType || totalCols <= 0) + return op.emitOpError("failed to infer tmrgsort format2 tmp type"); + pto::TileBufType tmpType = + makeVecTmpType(ctx, {1, totalCols}, elementType, {1, totalCols}); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), tmpType); + if (failed(tmp)) + return failure(); + OperationState state(op.getLoc(), op->getName()); + state.addOperands(operands); + state.addOperands(op.getDsts()); + state.addOperands(*tmp); + state.addOperands(op.getExcuted()); + state.addAttribute( + "operandSegmentSizes", + builder.getDenseI32ArrayAttr( + {static_cast(op.getSrcs().size()), 0, 1, 1, 1})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + struct PTOMaterializeImplicitTmpPass : public PassWrapper> { @@ -274,6 +879,86 @@ struct PTOMaterializeImplicitTmpPass failed = true; } + SmallVector optionalTmpOps; + func.walk([&](Operation *op) { + if (isa(op)) + optionalTmpOps.push_back(op); + }); + + for (Operation *op : optionalTmpOps) { + LogicalResult result = + llvm::TypeSwitch(op) + .Case([&](auto typedOp) { + return replaceTColSumWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTGatherWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTQuantWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTPowWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTPowSWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTSort32WithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTXorWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTXorSWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return materializeTCvtTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return materializeTMrgSortTmp(typedOp, requireExplicitTmp, ctx); + }) + .Default([](Operation *) { return success(); }); + if (mlir::failed(result)) + failed = true; + } + + SmallVector rowReductionOps; + func.walk([&](Operation *op) { + if (isa(op)) + rowReductionOps.push_back(op); + }); + + for (Operation *op : rowReductionOps) { + LogicalResult result = + llvm::TypeSwitch(op) + .Case([&](auto typedOp) { + return replaceRowReductionWithTmp(typedOp, requireExplicitTmp, + ctx); + }) + .Default([](Operation *) { return success(); }); + if (mlir::failed(result)) + failed = true; + } + + SmallVector mandatoryTmpOps; + func.walk([&](Operation *op) { + if (isa(op)) + mandatoryTmpOps.push_back(op); + }); + for (Operation *op : mandatoryTmpOps) { + if (mlir::failed( + materializeFixedMandatoryTmp(op, requireExplicitTmp, ctx))) + failed = true; + } + if (failed) signalPassFailure(); } diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 6a92462fd4..1f669c556a 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -586,8 +586,14 @@ struct PlannerAnalysis { if (outputRoots.empty()) return; - for (Value scratch : getWrittenNonDpsOperands(op, dpsInits)) + for (Value scratch : getWrittenNonDpsOperands(op, dpsInits)) { addForbidAliasBetweenRoots(getRoots(scratch), outputRoots); + for (Value operand : op->getOperands()) { + if (operand == scratch || llvm::is_contained(dpsInits, operand)) + continue; + addForbidAliasBetweenRoots(getRoots(scratch), getRoots(operand)); + } + } } void recordInplacePolicyConflicts(Operation *op, ValueRange dpsInits) { diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 271f9e50a8..e0aedbd0b8 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -9386,7 +9386,11 @@ struct PTOCvtToEmitC : public OpConversionPattern { Value satModeVal = rewriter.create( loc, satModeTy, emitc::OpaqueAttr::get(ctx, satTok)); - SmallVector operands{dst, src, rmodeVal, satModeVal}; + SmallVector operands{dst, src}; + if (adaptor.getTmp()) + operands.push_back(peelUnrealized(adaptor.getTmp())); + operands.push_back(rmodeVal); + operands.push_back(satModeVal); rewriter.create( loc, TypeRange{}, "TCVT", diff --git a/test/lit/pto/implicit_tmp_arg_reductions.pto b/test/lit/pto/implicit_tmp_arg_reductions.pto new file mode 100644 index 0000000000..4f51f1da7c --- /dev/null +++ b/test/lit/pto/implicit_tmp_arg_reductions.pto @@ -0,0 +1,31 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR + +module { + func.func @implicit_arg_reduction_tmps() { + %src = pto.alloc_tile : !pto.tile_buf + %col_max = pto.alloc_tile : !pto.tile_buf + %col_min = pto.alloc_tile : !pto.tile_buf + %row_max = pto.alloc_tile : !pto.tile_buf + %row_min = pto.alloc_tile : !pto.tile_buf + pto.tcolargmax ins(%src : !pto.tile_buf) outs(%col_max : !pto.tile_buf) + pto.tcolargmin ins(%src : !pto.tile_buf) outs(%col_min : !pto.tile_buf) + pto.trowargmax ins(%src : !pto.tile_buf) outs(%row_max : !pto.tile_buf) + pto.trowargmin ins(%src : !pto.tile_buf) outs(%row_min : !pto.tile_buf) + return + } +} + +// IR-LABEL: func.func @implicit_arg_reduction_tmps +// IR: pto.tcolargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// IR: pto.tcolargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// IR: pto.trowargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// IR: pto.trowargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) diff --git a/test/lit/pto/implicit_tmp_optional_level3_invalid.pto b/test/lit/pto/implicit_tmp_optional_level3_invalid.pto new file mode 100644 index 0000000000..0f5951cda9 --- /dev/null +++ b/test/lit/pto/implicit_tmp_optional_level3_invalid.pto @@ -0,0 +1,24 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --pto-level=level3 %s 2>&1 | FileCheck %s + +module { + func.func @implicit_tpow_tmp_level3() { + %addr0 = arith.constant 0 : i64 + %addr1 = arith.constant 256 : i64 + %addr2 = arith.constant 512 : i64 + %base = pto.alloc_tile addr = %addr0 : !pto.tile_buf + %exp = pto.alloc_tile addr = %addr1 : !pto.tile_buf + %dst = pto.alloc_tile addr = %addr2 : !pto.tile_buf + // CHECK: error: 'pto.tpow' op requires explicit tmp when PlanMemory is skipped + pto.tpow ins(%base, %exp : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/pto/implicit_tmp_optional_ops_materialization.pto b/test/lit/pto/implicit_tmp_optional_ops_materialization.pto new file mode 100644 index 0000000000..7660a4bc88 --- /dev/null +++ b/test/lit/pto/implicit_tmp_optional_ops_materialization.pto @@ -0,0 +1,80 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 + +module { + func.func @implicit_tcolsum_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tcolsum ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) {isBinary = true} + return + } + + func.func @implicit_tgather_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %indices = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tgather ins(%src, %indices : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @implicit_tquant_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tquant ins(%src, %fp : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) {quant_type = #pto} + return + } + + func.func @implicit_tpows_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + %scalar = arith.constant 2.0 : f32 + pto.tpows ins(%src, %scalar : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + return + } + + func.func @implicit_tsort32_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %idx = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tsort32 ins(%src, %idx : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK-LABEL: func.func @implicit_tcolsum_tmp +// CHECK: pto.alloc_tile addr = {{.*}} : !pto.tile_buf +// CHECK: pto.tcolsum ins(%{{.*}}, %{{.*}} {{.*}}isBinary = true + +// CHECK-LABEL: func.func @implicit_tgather_tmp +// CHECK: pto.tgather ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + +// CHECK-LABEL: func.func @implicit_tquant_tmp +// CHECK: pto.tquant ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) outs(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) + +// CHECK-LABEL: func.func @implicit_tpows_tmp +// CHECK: pto.tpows ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, f32, !pto.tile_buf) + +// CHECK-LABEL: func.func @implicit_tsort32_tmp +// CHECK: pto.tsort32 ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + +// A5-LABEL: func.func @implicit_tgather_tmp +// A5: pto.tgather ins(%{{.*}}, %{{.*}}, %{{.*}} : +// A5-LABEL: func.func @implicit_tquant_tmp +// A5: pto.tquant ins(%{{.*}}, %{{.*}} : {{.*}}) outs(%{{.*}} : +// A5-NOT: pto.tquant ins({{.*}}) outs(%{{.*}}, %{{.*}} +// A5-LABEL: func.func @implicit_tpows_tmp +// A5: pto.tpows ins(%{{.*}}, %{{.*}}, %{{.*}} : diff --git a/test/lit/pto/implicit_tmp_remaining_level3_invalid.pto b/test/lit/pto/implicit_tmp_remaining_level3_invalid.pto new file mode 100644 index 0000000000..50666f4310 --- /dev/null +++ b/test/lit/pto/implicit_tmp_remaining_level3_invalid.pto @@ -0,0 +1,46 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --pto-level=level3 %s 2>&1 | FileCheck %s + +module { + func.func @tprelu_missing_tmp() { + %a0 = arith.constant 0 : i64 + %a1 = arith.constant 256 : i64 + %a2 = arith.constant 512 : i64 + %src0 = pto.alloc_tile addr = %a0 : !pto.tile_buf + %src1 = pto.alloc_tile addr = %a1 : !pto.tile_buf + %dst = pto.alloc_tile addr = %a2 : !pto.tile_buf + pto.tprelu ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + + func.func @tcvt_missing_tmp() { + %a0 = arith.constant 0 : i64 + %a1 = arith.constant 256 : i64 + %src = pto.alloc_tile addr = %a0 : !pto.tile_buf + %dst = pto.alloc_tile addr = %a1 : !pto.tile_buf + pto.tcvt ins(%src {satmode = #pto} : !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + + func.func @tmrgsort_missing_tmp(%executed : vector<4xi16>) { + %a0 = arith.constant 0 : i64 + %a1 = arith.constant 256 : i64 + %a2 = arith.constant 512 : i64 + %src0 = pto.alloc_tile addr = %a0 : !pto.tile_buf + %src1 = pto.alloc_tile addr = %a1 : !pto.tile_buf + %dst = pto.alloc_tile addr = %a2 : !pto.tile_buf + pto.tmrgsort ins(%src0, %src1 no_tmp {exhausted = false} : !pto.tile_buf, !pto.tile_buf) outs(%dst, %executed : !pto.tile_buf, vector<4xi16>) + return + } +} + +// CHECK: error: 'pto.tprelu' op requires explicit tmp when PlanMemory is skipped +// CHECK: error: 'pto.tcvt' op requires explicit tmp for non-saturating narrowing tcvt when PlanMemory is skipped +// CHECK: error: 'pto.tmrgsort' op requires explicit tmp for tmrgsort format2 when PlanMemory is skipped diff --git a/test/lit/pto/implicit_tmp_remaining_ops_materialization.pto b/test/lit/pto/implicit_tmp_remaining_ops_materialization.pto new file mode 100644 index 0000000000..4f7aecd75e --- /dev/null +++ b/test/lit/pto/implicit_tmp_remaining_ops_materialization.pto @@ -0,0 +1,42 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --plan-memory-impl=modern --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @implicit_remaining_tmps(%executed : vector<4xi16>) { + %f32a = pto.alloc_tile : !pto.tile_buf + %f32b = pto.alloc_tile : !pto.tile_buf + %f32c = pto.alloc_tile : !pto.tile_buf + %mask = pto.alloc_tile : !pto.tile_buf + %i16dst = pto.alloc_tile : !pto.tile_buf + %scalar = arith.constant 3.0 : f32 + pto.tprelu ins(%f32a, %f32b : !pto.tile_buf, !pto.tile_buf) outs(%f32c : !pto.tile_buf) + pto.trem ins(%f32a, %f32b : !pto.tile_buf, !pto.tile_buf) outs(%f32c : !pto.tile_buf) + pto.trems ins(%f32a, %scalar : !pto.tile_buf, f32) outs(%f32c : !pto.tile_buf) + pto.tsel ins(%mask, %f32a, %f32b : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%f32c : !pto.tile_buf) + pto.tsels ins(%mask, %f32a, %scalar : !pto.tile_buf, !pto.tile_buf, f32) outs(%f32c : !pto.tile_buf) + pto.ttrans ins(%f32a : !pto.tile_buf) outs(%f32c : !pto.tile_buf) + pto.tcvt ins(%f32a {satmode = #pto} : !pto.tile_buf) outs(%i16dst : !pto.tile_buf) + %sort0 = pto.alloc_tile : !pto.tile_buf + %sort1 = pto.alloc_tile : !pto.tile_buf + %sortdst = pto.alloc_tile : !pto.tile_buf + pto.tmrgsort ins(%sort0, %sort1 no_tmp {exhausted = false} : !pto.tile_buf, !pto.tile_buf) outs(%sortdst, %executed : !pto.tile_buf, vector<4xi16>) + return + } +} + +// CHECK: pto.tprelu ins(%{{.*}}, %{{.*}}, %{{.*}} +// CHECK: pto.trem ins(%{{.*}}, %{{.*}}, %{{.*}} +// CHECK: pto.trems ins(%{{.*}}, %{{.*}}, %{{.*}} +// CHECK: pto.tsel ins(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} +// CHECK: pto.tsels ins(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} +// CHECK: pto.ttrans ins(%{{.*}}, %{{.*}} +// CHECK: pto.tcvt ins(%{{.*}}, %{{.*}} +// CHECK: pto.tmrgsort ins(%{{.*}}, %{{.*}}, %{{.*}} {exhausted = false} diff --git a/test/lit/pto/implicit_tmp_row_reductions.pto b/test/lit/pto/implicit_tmp_row_reductions.pto new file mode 100644 index 0000000000..cfd3ade699 --- /dev/null +++ b/test/lit/pto/implicit_tmp_row_reductions.pto @@ -0,0 +1,31 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR + +module { + func.func @implicit_row_reduction_tmps() { + %src = pto.alloc_tile : !pto.tile_buf + %max = pto.alloc_tile : !pto.tile_buf + %min = pto.alloc_tile : !pto.tile_buf + %sum = pto.alloc_tile : !pto.tile_buf + %prod = pto.alloc_tile : !pto.tile_buf + pto.trowmax ins(%src : !pto.tile_buf) outs(%max : !pto.tile_buf) + pto.trowmin ins(%src : !pto.tile_buf) outs(%min : !pto.tile_buf) + pto.trowsum ins(%src : !pto.tile_buf) outs(%sum : !pto.tile_buf) + pto.trowprod ins(%src : !pto.tile_buf) outs(%prod : !pto.tile_buf) + return + } +} + +// IR-LABEL: func.func @implicit_row_reduction_tmps +// IR: pto.trowmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// IR: pto.trowmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// IR: pto.trowsum ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// IR: pto.trowprod ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) diff --git a/test/lit/pto/implicit_tmp_xor_materialization.pto b/test/lit/pto/implicit_tmp_xor_materialization.pto new file mode 100644 index 0000000000..172d428c80 --- /dev/null +++ b/test/lit/pto/implicit_tmp_xor_materialization.pto @@ -0,0 +1,27 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR + +module { + func.func @implicit_xor_tmps() { + %scalar = arith.constant 3 : i16 + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %dst0 = pto.alloc_tile : !pto.tile_buf + %dst1 = pto.alloc_tile : !pto.tile_buf + pto.txor ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst0 : !pto.tile_buf) + pto.txors ins(%src0, %scalar : !pto.tile_buf, i16) outs(%dst1 : !pto.tile_buf) + return + } +} + +// IR-LABEL: func.func @implicit_xor_tmps +// IR: pto.txor ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) +// IR: pto.txors ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, i16, !pto.tile_buf) diff --git a/test/lit/pto/issue533_loop_zero_trip_sync_regression.pto b/test/lit/pto/issue533_loop_zero_trip_sync_regression.pto index 58e1c7b739..b213736de1 100644 --- a/test/lit/pto/issue533_loop_zero_trip_sync_regression.pto +++ b/test/lit/pto/issue533_loop_zero_trip_sync_regression.pto @@ -24,6 +24,7 @@ module attributes {pto.target_arch = "a2a3"} { %c128_i64 = arith.constant 128 : i64 %c160_i64 = arith.constant 160 : i64 %c4256_i64 = arith.constant 4256 : i64 + %c32768_i64 = arith.constant 32768 : i64 %c1024_index = arith.constant 1024 : index %c128_index = arith.constant 128 : index %c1_index = arith.constant 1 : index @@ -32,6 +33,7 @@ module attributes {pto.target_arch = "a2a3"} { %c0_index = arith.constant 0 : index %c8_index = arith.constant 8 : index %c16_index = arith.constant 16 : index + %rowexpand_tmp = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf %all_oi_tmp__co_l0_rv_v1_view = pto.make_tensor_view %arg0, shape = [%c1024_index, %c128_index], strides = [%c128_index, %c1_index] {layout = #pto.layout}: !pto.tensor_view %all_cur_mi__co_l0_rv_v1_view = pto.make_tensor_view %arg1, shape = [%c512_index, %c1_index], strides = [%c1_index, %c512_index] {layout = #pto.layout}: !pto.tensor_view %all_cur_li__co_l0_rv_v1_view = pto.make_tensor_view %arg2, shape = [%c512_index, %c1_index], strides = [%c1_index, %c512_index] {layout = #pto.layout}: !pto.tensor_view @@ -97,9 +99,9 @@ module attributes {pto.target_arch = "a2a3"} { pto.tadd ins(%li__rm_a0_tmp_v19, %li__rm_a1_tmp_v20 : !pto.tile_buf, !pto.tile_buf) outs(%li__row_major_tmp_v21 : !pto.tile_buf) %3 = pto.alloc_tile addr = %c96_i64 : !pto.tile_buf %4 = pto.alloc_tile addr = %c160_i64 : !pto.tile_buf - pto.trowexpandmul ins(%oi__tile, %alpha__tile : !pto.tile_buf, !pto.tile_buf) outs(%4 : !pto.tile_buf) + pto.trowexpandmul ins(%oi__tile, %alpha__tile, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%4 : !pto.tile_buf) %5 = pto.alloc_tile addr = %c10464_i64 : !pto.tile_buf - pto.trowexpandmul ins(%oi_tmp_valid__tile, %beta__tile : !pto.tile_buf, !pto.tile_buf) outs(%5 : !pto.tile_buf) + pto.trowexpandmul ins(%oi_tmp_valid__tile, %beta__tile, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%5 : !pto.tile_buf) %6 = pto.alloc_tile addr = %c10464_i64 : !pto.tile_buf pto.tadd ins(%4, %5 : !pto.tile_buf, !pto.tile_buf) outs(%6 : !pto.tile_buf) %mi__ssa_v3 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf @@ -111,7 +113,7 @@ module attributes {pto.target_arch = "a2a3"} { pto.tmov ins(%6 : !pto.tile_buf) outs(%oi__tile_mv : !pto.tile_buf) } %ctx__tile = pto.alloc_tile addr = %c6304_i64 : !pto.tile_buf - pto.trowexpanddiv ins(%oi__tile, %li__tile : !pto.tile_buf, !pto.tile_buf) outs(%ctx__tile : !pto.tile_buf) + pto.trowexpanddiv ins(%oi__tile, %li__tile, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%ctx__tile : !pto.tile_buf) %ctx_flat__tile = pto.alloc_tile addr = %c6304_i64 : !pto.tile_buf %ctx_flat_bf16__tile = pto.alloc_tile addr = %c4256_i64 : !pto.tile_buf pto.tcvt ins(%ctx_flat__tile{rmode = #pto} : !pto.tile_buf) outs(%ctx_flat_bf16__tile : !pto.tile_buf) diff --git a/test/lit/pto/issue533_loop_zero_trip_sync_regression_gss.pto b/test/lit/pto/issue533_loop_zero_trip_sync_regression_gss.pto index 56e75b410b..e8790c460f 100644 --- a/test/lit/pto/issue533_loop_zero_trip_sync_regression_gss.pto +++ b/test/lit/pto/issue533_loop_zero_trip_sync_regression_gss.pto @@ -24,6 +24,7 @@ module attributes {pto.target_arch = "a2a3"} { %c128_i64 = arith.constant 128 : i64 %c160_i64 = arith.constant 160 : i64 %c4256_i64 = arith.constant 4256 : i64 + %c32768_i64 = arith.constant 32768 : i64 %c1024_index = arith.constant 1024 : index %c128_index = arith.constant 128 : index %c1_index = arith.constant 1 : index @@ -32,6 +33,7 @@ module attributes {pto.target_arch = "a2a3"} { %c0_index = arith.constant 0 : index %c8_index = arith.constant 8 : index %c16_index = arith.constant 16 : index + %rowexpand_tmp = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf %all_oi_tmp__co_l0_rv_v1_view = pto.make_tensor_view %arg0, shape = [%c1024_index, %c128_index], strides = [%c128_index, %c1_index] {layout = #pto.layout}: !pto.tensor_view %all_cur_mi__co_l0_rv_v1_view = pto.make_tensor_view %arg1, shape = [%c512_index, %c1_index], strides = [%c1_index, %c512_index] {layout = #pto.layout}: !pto.tensor_view %all_cur_li__co_l0_rv_v1_view = pto.make_tensor_view %arg2, shape = [%c512_index, %c1_index], strides = [%c1_index, %c512_index] {layout = #pto.layout}: !pto.tensor_view @@ -97,9 +99,9 @@ module attributes {pto.target_arch = "a2a3"} { pto.tadd ins(%li__rm_a0_tmp_v19, %li__rm_a1_tmp_v20 : !pto.tile_buf, !pto.tile_buf) outs(%li__row_major_tmp_v21 : !pto.tile_buf) %3 = pto.alloc_tile addr = %c96_i64 : !pto.tile_buf %4 = pto.alloc_tile addr = %c160_i64 : !pto.tile_buf - pto.trowexpandmul ins(%oi__tile, %alpha__tile : !pto.tile_buf, !pto.tile_buf) outs(%4 : !pto.tile_buf) + pto.trowexpandmul ins(%oi__tile, %alpha__tile, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%4 : !pto.tile_buf) %5 = pto.alloc_tile addr = %c10464_i64 : !pto.tile_buf - pto.trowexpandmul ins(%oi_tmp_valid__tile, %beta__tile : !pto.tile_buf, !pto.tile_buf) outs(%5 : !pto.tile_buf) + pto.trowexpandmul ins(%oi_tmp_valid__tile, %beta__tile, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%5 : !pto.tile_buf) %6 = pto.alloc_tile addr = %c10464_i64 : !pto.tile_buf pto.tadd ins(%4, %5 : !pto.tile_buf, !pto.tile_buf) outs(%6 : !pto.tile_buf) %mi__ssa_v3 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf @@ -111,7 +113,7 @@ module attributes {pto.target_arch = "a2a3"} { pto.tmov ins(%6 : !pto.tile_buf) outs(%oi__tile_mv : !pto.tile_buf) } %ctx__tile = pto.alloc_tile addr = %c6304_i64 : !pto.tile_buf - pto.trowexpanddiv ins(%oi__tile, %li__tile : !pto.tile_buf, !pto.tile_buf) outs(%ctx__tile : !pto.tile_buf) + pto.trowexpanddiv ins(%oi__tile, %li__tile, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%ctx__tile : !pto.tile_buf) %ctx_flat__tile = pto.alloc_tile addr = %c6304_i64 : !pto.tile_buf %ctx_flat_bf16__tile = pto.alloc_tile addr = %c4256_i64 : !pto.tile_buf pto.tcvt ins(%ctx_flat__tile{rmode = #pto} : !pto.tile_buf) outs(%ctx_flat_bf16__tile : !pto.tile_buf) diff --git a/test/lit/pto/issue646_pipev_repeat_prune.pto b/test/lit/pto/issue646_pipev_repeat_prune.pto index d610b3326b..bb3584872b 100644 --- a/test/lit/pto/issue646_pipev_repeat_prune.pto +++ b/test/lit/pto/issue646_pipev_repeat_prune.pto @@ -9,11 +9,13 @@ module { %c0_i64 = arith.constant 0 : i64 %c4096_i64 = arith.constant 4096 : i64 %c8192_i64 = arith.constant 8192 : i64 + %c32768_i64 = arith.constant 32768 : i64 %acc = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf %src1 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf %scale = pto.alloc_tile addr = %c8192_i64 : !pto.tile_buf + %rowexpand_tmp = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf pto.tadd ins(%acc, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) - pto.trowexpanddiv ins(%acc, %scale : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) + pto.trowexpanddiv ins(%acc, %scale, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) return } @@ -25,11 +27,13 @@ module { %c0_i64 = arith.constant 0 : i64 %c4096_i64 = arith.constant 4096 : i64 %c12288_i64 = arith.constant 12288 : i64 + %c32768_i64 = arith.constant 32768 : i64 %acc = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf %src1 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf %scale = pto.alloc_tile addr = %c12288_i64 : !pto.tile_buf + %rowexpand_tmp = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf pto.tadd ins(%acc, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) - pto.trowexpanddiv ins(%acc, %scale : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) + pto.trowexpanddiv ins(%acc, %scale, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) return } @@ -41,11 +45,13 @@ module { %c0_i64 = arith.constant 0 : i64 %c4096_i64 = arith.constant 4096 : i64 %c8192_i64 = arith.constant 8192 : i64 + %c32768_i64 = arith.constant 32768 : i64 %acc = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf %src1 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf %scale = pto.alloc_tile addr = %c8192_i64 : !pto.tile_buf + %rowexpand_tmp = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf pto.tadd ins(%acc, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) - pto.trowexpanddiv ins(%acc, %scale : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) + pto.trowexpanddiv ins(%acc, %scale, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) return } @@ -60,14 +66,16 @@ module { %c8192_i64 = arith.constant 8192 : i64 %c12288_i64 = arith.constant 12288 : i64 %c16384_i64 = arith.constant 16384 : i64 + %c32768_i64 = arith.constant 32768 : i64 %acc = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf %src1 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf %scale = pto.alloc_tile addr = %c8192_i64 : !pto.tile_buf %tmp = pto.alloc_tile addr = %c12288_i64 : !pto.tile_buf %tmp_src = pto.alloc_tile addr = %c16384_i64 : !pto.tile_buf + %rowexpand_tmp = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf pto.tadd ins(%acc, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) pto.tadd ins(%tmp, %tmp_src : !pto.tile_buf, !pto.tile_buf) outs(%tmp : !pto.tile_buf) - pto.trowexpanddiv ins(%acc, %scale : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) + pto.trowexpanddiv ins(%acc, %scale, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) return } diff --git a/test/lit/pto/tpow_fp_missing_tmp_invalid.pto b/test/lit/pto/tpow_fp_missing_tmp_invalid.pto index 23d832af96..01813ab8ad 100644 --- a/test/lit/pto/tpow_fp_missing_tmp_invalid.pto +++ b/test/lit/pto/tpow_fp_missing_tmp_invalid.pto @@ -6,8 +6,9 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// Floating-point tpow requires tmp scratch (used by PowF / TPowFloat). -// RUN: not ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s +// Floating-point tpow gets an implicit tmp before memplan. +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 module { func.func @tpow_fp_missing_tmp() { @@ -19,4 +20,8 @@ module { } } -// CHECK: expects tmp when element type is floating-point +// A3: pto.alloc_tile addr = {{.*}} : !pto.tile_buf +// A3: pto.tpow ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + +// A5: pto.alloc_tile addr = {{.*}} : !pto.tile_buf +// A5: pto.tpow ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) From ebab68c69afc47b126377b60dc3028fbf5400847 Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 4 Aug 2026 16:23:33 +0800 Subject: [PATCH 045/122] =?UTF-8?q?Fix=20rowexpand=5Ftile=5Fnative.pto=20R?= =?UTF-8?q?UN=20level=20(level3=20=E2=86=92=20level2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This test verifies tile-native IR preservation (!pto.tile_buf kept native, not lowered to memref), which is a level2 behavior. Running it at level3 invokes the full pipeline including PlanMemory skip, which triggers the new implicit-tmp verifier rejecting mode-1 col-major row-expand ops without explicit tmp (since implicit tmp cannot be materialized when PlanMemory is skipped at level3). At level2, PlanMemory runs and materializes implicit tmp for mode-1 row-expand ops correctly, matching the test's tile-native intent. --- test/lit/pto/rowexpand_tile_native.pto | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/lit/pto/rowexpand_tile_native.pto b/test/lit/pto/rowexpand_tile_native.pto index 61d9f7a4ec..a7523bae6b 100644 --- a/test/lit/pto/rowexpand_tile_native.pto +++ b/test/lit/pto/rowexpand_tile_native.pto @@ -6,8 +6,8 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE -// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC +// RUN: ptoas --pto-level=level2 --pto-arch=a5 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level2 --pto-arch=a3 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC module { func.func private @trowexpand_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { From 55a3f5e7ad8f3d21cc16347db4059e9811ae32c8 Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 4 Aug 2026 17:18:17 +0800 Subject: [PATCH 046/122] Fix implicit-tmp tile-native lit tests after rebase onto main After rebasing onto main (which removed memref compatibility and added the optional tmp operand to tcvt), four lit tests had stale inputs or CHECK lines that no longer matched the generated IR: - tci_implicit_tmp_materialization.pto: rewrite inputs from memref to tile-native partition_tensor_view so the test parses on main. - tquant_no_implicit_tmp_a3.pto: this is a negative test expecting the implicit-tmp verifier to reject a dynamic-valid-shape tquant src. Use 'not ptoas' (the repo's standard negative-test idiom) so the pipe exit code reflects the expected failure, and match the emitted diagnostic instead of an unreachable func label. - cvt_tile_native.pto / tcvt_low_precision_a5_valid.pto: the optional tmp operand now prints 'operandSegmentSizes = array' in the tcvt IR; wildcard it in the CHECK lines. Verified on remote A3 (LLVM21): full check-pto = 1565 passed, 0 failed, 1 unsupported. --- test/lit/pto/cvt_tile_native.pto | 2 +- .../pto/tci_implicit_tmp_materialization.pto | 19 +++++++++++++++---- test/lit/pto/tcvt_low_precision_a5_valid.pto | 6 +++--- test/lit/pto/tquant_no_implicit_tmp_a3.pto | 18 +++++++++--------- 4 files changed, 28 insertions(+), 17 deletions(-) diff --git a/test/lit/pto/cvt_tile_native.pto b/test/lit/pto/cvt_tile_native.pto index af37bd6e13..c33ff11029 100644 --- a/test/lit/pto/cvt_tile_native.pto +++ b/test/lit/pto/cvt_tile_native.pto @@ -23,7 +23,7 @@ module { } // NATIVE-LABEL: func.func private @tcvt_arg( -// NATIVE: pto.tcvt ins(%arg0 {rmode = #pto, satmode = #pto} : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) +// NATIVE: pto.tcvt ins(%arg0 {{.*}}rmode = #pto, satmode = #pto{{.*}} : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) // NATIVE-NOT: memref< // EMITC-LABEL: tcvt_arg( diff --git a/test/lit/pto/tci_implicit_tmp_materialization.pto b/test/lit/pto/tci_implicit_tmp_materialization.pto index 78cec449e5..e8fbd41bcb 100644 --- a/test/lit/pto/tci_implicit_tmp_materialization.pto +++ b/test/lit/pto/tci_implicit_tmp_materialization.pto @@ -1,16 +1,27 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS PROGRAM IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// TCI without an explicit tmp operand: ptoas should materialize an implicit +// tmp (a 1x192 fp32 buffer on A3) and pass it to the PTO-ISA TCI interface, +// then memory-plan all variables together with the tmp. + // RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR // RUN: ptoas --pto-arch=a3 --pto-level=level2 %s 2>&1 | FileCheck %s --check-prefix=CPP // RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5IR module { - func.func @tci_implicit_tmp(%dst: memref<32xi32, #pto.address_space>) { + func.func @tci_implicit_tmp(%dst: !pto.partition_tensor_view<1x32xi32>) { %c0_i32 = arith.constant 0 : i32 - %src = memref.reinterpret_cast %dst to offset: [0], sizes: [1, 32], strides: [32, 1] {layout = #pto.layout} : memref<32xi32, #pto.address_space> to memref<1x32xi32, strided<[32, 1], offset: ?>, #pto.address_space> - %tile = pto.alloc_tile : !pto.tile_buf + %tile = pto.declare_tile -> !pto.tile_buf pto.tci ins(%c0_i32 : i32) outs(%tile : !pto.tile_buf) pto.tstore ins(%tile : !pto.tile_buf) - outs(%src : memref<1x32xi32, strided<[32, 1], offset: ?>, #pto.address_space>) {layout = #pto.layout, pto.inferred_layout = true} + outs(%dst : !pto.partition_tensor_view<1x32xi32>) return } } diff --git a/test/lit/pto/tcvt_low_precision_a5_valid.pto b/test/lit/pto/tcvt_low_precision_a5_valid.pto index 15cc41a6da..aa9a8271b1 100644 --- a/test/lit/pto/tcvt_low_precision_a5_valid.pto +++ b/test/lit/pto/tcvt_low_precision_a5_valid.pto @@ -30,9 +30,9 @@ module { // CHECK: func.func @tcvt_low_precision_a5_valid() attributes {pto.kernel_kind = #pto.kernel_kind} // CHECK: pto.declare_tile -> !pto.tile_buf // CHECK: pto.declare_tile -> !pto.tile_buf -// CHECK: pto.tcvt ins(%{{.*}} {rmode = #pto, satmode = #pto} : !pto.tile_buf +// CHECK: pto.tcvt ins(%{{.*}} {{.*}}rmode = #pto, satmode = #pto{{.*}} : !pto.tile_buf // CHECK: outs(%{{.*}} : !pto.tile_buf) -// CHECK: pto.tcvt ins(%{{.*}} {rmode = #pto, satmode = #pto} : !pto.tile_buf +// CHECK: pto.tcvt ins(%{{.*}} {{.*}}rmode = #pto, satmode = #pto{{.*}} : !pto.tile_buf // CHECK: outs(%{{.*}} : !pto.tile_buf) -// CHECK: pto.tcvt ins(%{{.*}} {rmode = #pto, satmode = #pto} : !pto.tile_buf +// CHECK: pto.tcvt ins(%{{.*}} {{.*}}rmode = #pto, satmode = #pto{{.*}} : !pto.tile_buf // CHECK: outs(%{{.*}} : !pto.tile_buf) diff --git a/test/lit/pto/tquant_no_implicit_tmp_a3.pto b/test/lit/pto/tquant_no_implicit_tmp_a3.pto index be39a34079..02c5ba0d8d 100644 --- a/test/lit/pto/tquant_no_implicit_tmp_a3.pto +++ b/test/lit/pto/tquant_no_implicit_tmp_a3.pto @@ -2,11 +2,16 @@ // This program is free software, you can redistribute it and/or modify it under the terms and conditions of // CANN Open Software License Agreement Version 2.0 (the "License"). // Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// THIS PROGRAM IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-arch=a3 %s -emit-pto-ir 2>&1 | FileCheck %s +// TQUANT with a dynamic-valid-shape src and no explicit tmp cannot have its +// implicit tmp materialized (the tmp type is derived from the static src +// shape). The implicit-tmp pass must reject it with a clear diagnostic +// instead of silently dropping the tmp. + +// RUN: not ptoas --pto-arch=a3 %s -emit-pto-ir 2>&1 | FileCheck %s module { func.func @tquant_no_implicit_tmp_a3(%valid_row: index, %valid_col: index) @@ -15,13 +20,6 @@ module { %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile valid_row = %valid_row valid_col = %valid_col : !pto.tile_buf - // CHECK-LABEL: func.func @tquant_no_implicit_tmp_a3 - // CHECK-SAME: (%[[ROW:arg[0-9]+]]: index, %[[COL:arg[0-9]+]]: index) - // CHECK: pto.alloc_tile{{.*}}valid_row = %[[ROW]] valid_col = %[[COL]]{{.*}} : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) @@ -29,3 +27,5 @@ module { return } } + +// CHECK: 'pto.tquant' op requires static tile_buf src to materialize implicit tquant tmp From 68e079d8c2c2188b37131ef702cc28562d1b889b Mon Sep 17 00:00:00 2001 From: FangRui Date: Wed, 5 Aug 2026 15:14:03 +0800 Subject: [PATCH 047/122] Fix ptodsl bindings and op printers for optional tmp operand Two CI failures after adding the optional tmp operand: 1. TypeError: trowmax() takes 2 positional arguments but 3 were given - The MLIR-generated Python bindings expose tmp as a keyword-only argument (tmp=None after *) for ops whose ODS declares Optional:$tmp. The ptodsl wrappers in _ops.py were still passing tmp positionally, which collided with the dst slot. - Fix trowsum/max/min/prod/argmax/argmin, tcolargmax/argmin, txor, txors to pass tmp=... as a keyword argument (matching tcolsum, tcvt, tsel, tci, tmrgsort, tgather which already did so). 2. docs_as_test cvt round-trip textual drift - After adding the optional tmp operand, MLIR auto-attaches an 'operandSegmentSizes' attribute to tcvt/tpow/tpows/tcolsum, which the custom assembly printers leaked via printOptionalAttrDict, breaking the parse->print round-trip stability checked by the docs_as_test harness. - Elide 'operandSegmentSizes' in TCvtOp/TPowOp/TPowSOp/TColSumOp printers (the repo's established idiom, already used by TRowExpand*/TXor/TGather/TSort32/TMrgSort). Verified on remote A3 (LLVM21): ptodsl_docs_as_test PASS, test_ptoas_frontend_verify PASS, full check-pto = 1565 passed, 0 failed, 1 unsupported. --- lib/PTO/IR/PTO.cpp | 13 ++++++++----- ptodsl/ptodsl/_ops.py | 20 ++++++++++---------- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 984e0ad7d4..5d98065c50 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -6337,7 +6337,7 @@ void mlir::pto::TColSumOp::print(OpAsmPrinter &p) { // Format 2: ins(%src, %tmp {isBinary = ...}: type, type) outs(%dst : type) p << " ins(" << getSrc() << ", " << getTmp(); // Print isBinary attribute if present - SmallVector elidedAttrs; + SmallVector elidedAttrs = {"operandSegmentSizes"}; if (!getIsBinaryAttr() || getIsBinaryAttr().getValue() == false) { elidedAttrs.push_back("isBinary"); } @@ -6352,7 +6352,7 @@ void mlir::pto::TColSumOp::print(OpAsmPrinter &p) { // Print remaining attributes for format 1 (excluding isBinary) if (!getTmp()) { - SmallVector elidedAttrs = {"isBinary"}; + SmallVector elidedAttrs = {"isBinary", "operandSegmentSizes"}; p.printOptionalAttrDict((*this)->getAttrs(), elidedAttrs); } } @@ -9645,7 +9645,8 @@ void mlir::pto::TCvtOp::print(OpAsmPrinter &p) { } attrs.set(attr.getName(), attr.getValue()); } - p.printOptionalAttrDict(attrs.getAttrs()); + p.printOptionalAttrDict(attrs.getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); p << " : " << getSrc().getType(); if (getTmp()) p << ", " << getTmp().getType(); @@ -11710,7 +11711,8 @@ void mlir::pto::TPowOp::print(OpAsmPrinter &p) { p << ", " << getTmp().getType(); p << ")"; p << " outs(" << getDst() << " : " << getDst().getType() << ")"; - p.printOptionalAttrDict((*this)->getAttrs()); + p.printOptionalAttrDict((*this)->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); } // TPOWS assembly format: @@ -11767,7 +11769,8 @@ void mlir::pto::TPowSOp::print(OpAsmPrinter &p) { p << ", " << getTmp().getType(); p << ")"; p << " outs(" << getDst() << " : " << getDst().getType() << ")"; - p.printOptionalAttrDict((*this)->getAttrs()); + p.printOptionalAttrDict((*this)->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); } static ParseResult parseTRowExpandBinaryLikeOp(OpAsmParser &parser, diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 914c39695b..5ec91c9a66 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -3490,8 +3490,8 @@ def trowsum(src, tmp, dst): """``pto.trowsum ins(src, tmp) outs(dst)``.""" _pto.trowsum( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3499,8 +3499,8 @@ def trowmax(src, tmp, dst): """``pto.trowmax ins(src, tmp) outs(dst)``.""" _pto.trowmax( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3508,8 +3508,8 @@ def trowmin(src, tmp, dst): """``pto.trowmin ins(src, tmp) outs(dst)``.""" _pto.trowmin( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3517,8 +3517,8 @@ def trowprod(src, tmp, dst): """``pto.trowprod ins(src, tmp) outs(dst)``.""" _pto.trowprod( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3526,8 +3526,8 @@ def trowargmax(src, tmp, dst): """``pto.trowargmax ins(src, tmp) outs(dst)``.""" _pto.trowargmax( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3535,8 +3535,8 @@ def trowargmin(src, tmp, dst): """``pto.trowargmin ins(src, tmp) outs(dst)``.""" _pto.trowargmin( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3578,8 +3578,8 @@ def tcolargmax(src, tmp, dst): """``pto.tcolargmax ins(src, tmp) outs(dst)``.""" _pto.tcolargmax( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3587,8 +3587,8 @@ def tcolargmin(src, tmp, dst): """``pto.tcolargmin ins(src, tmp) outs(dst)``.""" _pto.tcolargmin( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -4118,8 +4118,8 @@ def txor(src0, src1, tmp, dst): _pto.txor( unwrap_surface_value(src0), unwrap_surface_value(src1), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -4128,8 +4128,8 @@ def txors(src, scalar, tmp, dst): _pto.txors( unwrap_surface_value(src), _coerce_tile_scalar_operand(src, scalar, context="txors"), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) From 110a9992dbd8e9009a5158032ddbb3ac5a5cf3d3 Mon Sep 17 00:00:00 2001 From: FangRui Date: Wed, 5 Aug 2026 16:01:03 +0800 Subject: [PATCH 048/122] Resolve TGather verifier conflict after rebase onto main main added a stricter TGather contract (PR #1080 'Add TGATHER indices and mask'): A2/A3 index-form and all compare-form tgather ops now require an explicit tmp ("expects both indices and tmp"), while A5 index-form deliberately emits WITHOUT tmp (TGATHER(src, indices, dst)). This conflicts with the implicit-tmp design, which materialized a tmp for any omitted-tmp tgather. Restore main's contract and stop materializing tgather tmp: - PTO.cpp TGatherOp::verify: revert to main's logic (A2/A3 needs tmp, A5 index-form lets verifyIndexForm handle the missing tmp). Remove the implicit-tmp early-return 'if (!getTmp()) return success()'. - PTOMaterializeImplicitTmp.cpp: drop TGatherOp from the optional-tmp dispatch and isa<...> list (tgather no longer gets an implicit tmp). Mark replaceTGatherWithTmp [[maybe_unused]] to keep the documented implementation around without a -Werror unused-function failure. - implicit_tmp_optional_ops_materialization.pto: drop the implicit_tgather_tmp test case (A2/A3 index-form now requires an explicit tmp; the implicit-tmp path no longer applies to tgather). Verified on remote A3 (LLVM21): full check-pto = 1646 passed, 0 failed, 1 unsupported; ptodsl_docs_as_test + frontend_verify PASS. --- lib/PTO/IR/PTO.cpp | 20 ++++++------------- .../Transforms/PTOMaterializeImplicitTmp.cpp | 11 ++++------ ...licit_tmp_optional_ops_materialization.pto | 14 ------------- 3 files changed, 10 insertions(+), 35 deletions(-) diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 5d98065c50..202fefe1e3 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -7845,18 +7845,14 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() { if (getAxisAttr()) return emitOpError("axis attribute must not be provided without maskPattern"); if (getCdst() || getKValue()) { - if (!getCdst() || !getKValue()) - return emitOpError("compare-form tgather expects dst, cdst, and kValue"); + if (!getCdst() || !getKValue() || !getTmp()) + return emitOpError("compare-form tgather expects dst, cdst, kValue, and tmp"); if (getIndices()) return emitOpError("compare-form tgather does not take indices"); - if (!getTmp()) - return success(); return verifyCompareForm(/*allowA5SrcTypes=*/false); } - if (!getIndices()) - return emitOpError("index-form tgather expects indices"); - if (!getTmp()) - return success(); + if (!getIndices() || !getTmp()) + return emitOpError("index-form tgather expects both indices and tmp"); return verifyIndexForm(/*allow16BitIndices=*/false, /*allowA5ElemTypes=*/false); }; @@ -7869,18 +7865,14 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() { if (getAxisAttr()) return emitOpError("axis attribute must not be provided without maskPattern"); if (getCdst() || getKValue()) { - if (!getCdst() || !getKValue()) - return emitOpError("compare-form tgather expects dst, cdst, and kValue"); + if (!getCdst() || !getKValue() || !getTmp()) + return emitOpError("compare-form tgather expects dst, cdst, kValue, and tmp"); if (getIndices()) return emitOpError("compare-form tgather does not take indices"); - if (!getTmp()) - return success(); return verifyCompareForm(/*allowA5SrcTypes=*/true); } if (!getIndices()) return emitOpError("index-form tgather expects indices"); - if (!getTmp()) - return success(); return verifyIndexForm(/*allow16BitIndices=*/true, /*allowA5ElemTypes=*/true); }; diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp index 8c0d60eef7..d35c7147e3 100644 --- a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -376,9 +376,9 @@ static LogicalResult replaceTPowSWithTmp(pto::TPowSOp op, return success(); } -static LogicalResult replaceTGatherWithTmp(pto::TGatherOp op, - bool requireExplicitTmp, - MLIRContext *ctx) { +[[maybe_unused]] static LogicalResult +replaceTGatherWithTmp(pto::TGatherOp op, bool requireExplicitTmp, + MLIRContext *ctx) { if (op.getTmp() || op.hasMaskForm()) return success(); if (!op.hasIndexForm() && !op.hasCompareForm()) @@ -881,7 +881,7 @@ struct PTOMaterializeImplicitTmpPass SmallVector optionalTmpOps; func.walk([&](Operation *op) { - if (isa(op)) optionalTmpOps.push_back(op); @@ -893,9 +893,6 @@ struct PTOMaterializeImplicitTmpPass .Case([&](auto typedOp) { return replaceTColSumWithTmp(typedOp, requireExplicitTmp, ctx); }) - .Case([&](auto typedOp) { - return replaceTGatherWithTmp(typedOp, requireExplicitTmp, ctx); - }) .Case([&](auto typedOp) { return replaceTQuantWithTmp(typedOp, requireExplicitTmp, ctx); }) diff --git a/test/lit/pto/implicit_tmp_optional_ops_materialization.pto b/test/lit/pto/implicit_tmp_optional_ops_materialization.pto index 7660a4bc88..c7b4736c44 100644 --- a/test/lit/pto/implicit_tmp_optional_ops_materialization.pto +++ b/test/lit/pto/implicit_tmp_optional_ops_materialization.pto @@ -18,15 +18,6 @@ module { return } - func.func @implicit_tgather_tmp() { - %src = pto.alloc_tile : !pto.tile_buf - %indices = pto.alloc_tile : !pto.tile_buf - %dst = pto.alloc_tile : !pto.tile_buf - pto.tgather ins(%src, %indices : !pto.tile_buf, !pto.tile_buf) - outs(%dst : !pto.tile_buf) - return - } - func.func @implicit_tquant_tmp() { %src = pto.alloc_tile : !pto.tile_buf %fp = pto.alloc_tile : !pto.tile_buf @@ -59,9 +50,6 @@ module { // CHECK: pto.alloc_tile addr = {{.*}} : !pto.tile_buf // CHECK: pto.tcolsum ins(%{{.*}}, %{{.*}} {{.*}}isBinary = true -// CHECK-LABEL: func.func @implicit_tgather_tmp -// CHECK: pto.tgather ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) - // CHECK-LABEL: func.func @implicit_tquant_tmp // CHECK: pto.tquant ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) outs(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) @@ -71,8 +59,6 @@ module { // CHECK-LABEL: func.func @implicit_tsort32_tmp // CHECK: pto.tsort32 ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) -// A5-LABEL: func.func @implicit_tgather_tmp -// A5: pto.tgather ins(%{{.*}}, %{{.*}}, %{{.*}} : // A5-LABEL: func.func @implicit_tquant_tmp // A5: pto.tquant ins(%{{.*}}, %{{.*}} : {{.*}}) outs(%{{.*}} : // A5-NOT: pto.tquant ins({{.*}}) outs(%{{.*}}, %{{.*}} From 0c98de307194a931b41ec3890c227ab862b51adb Mon Sep 17 00:00:00 2001 From: FangRui Date: Wed, 5 Aug 2026 16:13:07 +0800 Subject: [PATCH 049/122] Update implicit-tmp design doc after main removed memref IR main removed the memref compatibility layer (PTOViewToMemref, PTOMaterializeTileHandles, memref.alloc, pto.pointer_cast / bind_tile ops) and tightened the TGather tmp contract (PR #1080). The design doc still referenced these, so refresh it: - Pipeline diagram: replace deleted PTOViewToMemref / PTOMaterializeTileHandles / PTOToEmitC with the actual current passes (PTOFusionRegionGen -> pto-materialize-implicit-tmp -> PTORematerializeFixpipeVectorQuant -> pto-plan-memory -> PTOResolveReservedBuffers -> sync passes -> PTOResolveBufferSelect -> EmitPTOManual). - Drop the 'memref.alloc' / 'pto.pointer_cast' / 'pto.bind_tile' bullets (those ops no longer exist) and the now-meaningless 'CHECK-NOT: memref.alloc' lit assertions; note that optional-tmp ops elide 'operandSegmentSizes' in the custom printer instead. - Rename stray PTOToEmitC references to EmitPTOManual (the real PTO->EmitC pass name). - TGATHER: mark it as NOT participating in implicit-tmp materialize. main requires an explicit tmp on A2/A3 index/compare-form and emits A5 index-form without tmp, so tgather is removed from the materialize dispatch; add a dedicated classification row and rewrite the TGATHER data-movement section accordingly. --- ...oas-implicit-tmp-materialization-design.md | 36 ++++++++++--------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/docs/designs/ptoas-implicit-tmp-materialization-design.md b/docs/designs/ptoas-implicit-tmp-materialization-design.md index 719563b370..c9c3c9cdec 100644 --- a/docs/designs/ptoas-implicit-tmp-materialization-design.md +++ b/docs/designs/ptoas-implicit-tmp-materialization-design.md @@ -33,16 +33,17 @@ PTOAS 前端 IR 中很多 tile op 的 `tmp` operand 是可选的。当前如果 pto-materialize-implicit-tmp ``` -该 pass 运行在 `PTOViewToMemref` 之后、`pto-plan-memory` 之前： +该 pass 运行在 fusion/调度 pass 之后、`pto-plan-memory` 之前： ```text -PTOViewToMemref +PTOFusionRegionGen -> pto-materialize-implicit-tmp - -> pto-plan-memory + -> PTORematerializeFixpipeVectorQuant + -> pto-plan-memory (level1/level2 only; skipped at level3) -> PTOResolveReservedBuffers - -> sync passes - -> PTOMaterializeTileHandles - -> PTOToEmitC + -> sync passes (InsertSync / GraphSyncSolver / BarrierAll ...) + -> PTOResolveBufferSelect + -> EmitPTOManual (PTO -> EmitC lowering) ``` pass 的职责是扫描所有已纳入改造的目标 op。如果 op 没有 tmp operand，就根据该 op 的 `TmpRequirement` 在 op 前插入 `pto.alloc_tile(no addr)`，并重写原 op，使其显式携带 tmp。 @@ -72,8 +73,6 @@ bool requireExplicitAtLevel3; 对自动生成的 tmp： - 使用 tile-native `pto.alloc_tile(no addr)`。 -- 不创建 `memref.alloc`。 -- 不创建 `pto.pointer_cast` / `pto.bind_tile`。 - 不设置 `addr`，由 memplan 统一规划。 - 尽量使用静态 full-valid shape，即 `v_row/v_col` 与 `rows/cols` 一致，不额外携带 `valid_row` / `valid_col` operand。 - 定义位置必须支配目标 op。 @@ -140,7 +139,7 @@ level3 + target_op(no tmp) => pass/verifier 报错 引入 `pto-materialize-implicit-tmp` 后，level1/level2 的目标 op 在 EmitC 前都会携带 tmp，因此会自然走带 tmp 的 overload。 -不建议在 PTOToEmitC 中补 tmp，原因： +不建议在 `EmitPTOManual`（PTO -> EmitC lowering）中补 tmp，原因： - EmitC 阶段已经错过 memplan。 - 临时生成 tmp 无法获得 local addr。 @@ -158,7 +157,7 @@ level3 + target_op(no tmp) => pass/verifier 报错 Optional:$tmp ``` -PTOToEmitC 也已经根据 `op.getTmp()` 选择带 tmp 或不带 tmp 的 C++ 调用。因此 TCI 改造不需要改 `pto.tci` 的 IR 语法，关键是保证进入 EmitC 前缺省 tmp 已经被显式 materialize。 +`EmitPTOManual`（PTO -> EmitC lowering）也已经根据 `op.getTmp()` 选择带 tmp 或不带 tmp 的 C++ 调用。因此 TCI 改造不需要改 `pto.tci` 的 IR 语法，关键是保证进入 EmitC 前缺省 tmp 已经被显式 materialize。 TCI 后端 C++ 接口存在两类 overload： @@ -326,9 +325,10 @@ test/lit/pto/tci_implicit_tmp_materialization.pto CHECK: pto.alloc_tile CHECK-SAME: dtype=f32 CHECK: pto.tci ins(%{{.*}}, %{{.*}} -CHECK-NOT: memref.alloc ``` +可选地检查 IR 不打印 `operandSegmentSizes`（可选 tmp 的自定义 assembly printer 已 elide 该属性，保证 round-trip 稳定）。 + 可以额外检查自动生成 shape：b32 dst 为 `f32 1x192`，b16 dst 为 `f32 1x448`。同时应覆盖 A5 下不自动生成 tmp。 #### lit：memplan 回写 addr @@ -527,7 +527,6 @@ forbidAlias(tmp, dst) ```text CHECK: pto.alloc_tile CHECK: pto.trowexpandadd ins(%{{.*}}, %{{.*}}, %{{.*}} -CHECK-NOT: memref.alloc ``` 同时检查 A5 下不自动生成 tmp。 @@ -584,11 +583,12 @@ TSORT32, TQUANT | 分类 | Op / 模式 | 设计结论 | | --- | --- | --- | -| A2/A3 使用 tmp，A5 接受但不使用 | `TCOLARGMAX`、`TCOLARGMIN`、`TROWARGMAX`、`TROWARGMIN`、`TGATHER`、`TROWPROD`、`TROWSUM`、`TROWMAX`、`TROWMIN`、`TSEL`、`TSELS`、`TREM`、`TREMS`、`TQUANT`、`TADDDEQRELU` | level1/2 在 A2/A3 生成真实 scratch；若 A5 C++ 签名仍要求 tmp，则生成不带 MemoryEffects 的 ABI placeholder。 | +| A2/A3 使用 tmp，A5 接受但不使用 | `TCOLARGMAX`、`TCOLARGMIN`、`TROWARGMAX`、`TROWARGMIN`、`TROWPROD`、`TROWSUM`、`TROWMAX`、`TROWMIN`、`TSEL`、`TSELS`、`TREM`、`TREMS`、`TQUANT`、`TADDDEQRELU` | level1/2 在 A2/A3 生成真实 scratch；若 A5 C++ 签名仍要求 tmp，则生成不带 MemoryEffects 的 ABI placeholder。 | | A2/A3 和 A5 都可能使用 tmp | `TCOLSUM(isBinary=true)`、`TSORT32` 非 32 对齐尾部、`TMRGSORT` 多列表归并 format2 | tmp 使用由 op 模式决定，不能只按 arch 判断。 | | 条件性 tmp，不应无条件 materialize | `TTRANS`、`TCVT`、`TPOW`、`TPOWS`、`TRSQRT`、`TMRGSORT`、`TSORT32` | 需要先判断精度、dtype、layout、format 或尾部条件。 | | 已从 mandatory tmp 改为 optional tmp | `TTRANS`、`TXOR`、`TXORS`、`TPRELU`、`TROWPROD`、`TROWSUM`、`TROWMAX`、`TROWMIN`、`TROWARGMAX`、`TROWARGMIN`、`TCOLARGMAX`、`TCOLARGMIN`、`TSEL`、`TSELS`、`TREM`、`TREMS` | ODS、parse/print、verifier、MemoryEffects、materialize 和 lowering 已接入。 | -| 当前 PTOAS IR 已有 optional tmp | `TGATHER`、`TCOLSUM`、`TRSQRT`、`TPOW`、`TPOWS`、`TSORT32`、`TQUANT` | 可直接纳入 `pto-materialize-implicit-tmp` 的后续实现。 | +| 当前 PTOAS IR 已有 optional tmp | `TCOLSUM`、`TRSQRT`、`TPOW`、`TPOWS`、`TSORT32`、`TQUANT` | 可直接纳入 `pto-materialize-implicit-tmp` 的后续实现。 | +| 已有 optional tmp 但不纳入 implicit-tmp materialize | `TGATHER` | main 分支要求 A2/A3 显式 tmp（verifier 拒绝省略），A5 index-form 设计为无 tmp；`replaceTGatherWithTmp` 实现保留但当前不在 dispatch 中启用。 | | 当前 PTOAS IR 暂无对应 op | `TADDDEQRELU` | 需先完成 PTOAS IR 接入；`TCVT` 已新增 optional tmp operand。 | ### 通用规则 @@ -766,11 +766,13 @@ TGATHER, TTRANS, TCVT TGATHER： -- 当前 PTOAS IR 已支持 optional tmp。 +- 当前 PTOAS IR 已支持 optional tmp，但 **tgather 不纳入 implicit-tmp materialize 范围**。 +- main 分支（PR #1080 "Add TGATHER indices and mask"）对 tgather 的 tmp 契约更严：A2/A3 index-form 和所有 compare-form 都要求显式 `tmp`（verifier 报 `index-form tgather expects both indices and tmp` / `compare-form tgather expects dst, cdst, kValue, and tmp`）；A5 index-form 设计为不带 tmp（emit `TGATHER(src, indices, dst)` 三参数）。 +- 因此 tgather 省略 tmp 时不由 `pto-materialize-implicit-tmp` 自动补齐，而是由 verifier 直接拒绝（A2/A3）或允许无 tmp（A5 index-form）。 - index form：A2/A3 C++ API 需要 tmp；tmp dtype 与 indices dtype 一致，shape 覆盖 indices；A5 不使用 tmp。 - compare form：A2/A3 tmp 是合并暂存缓冲区，包含 `cmpsTmp`、`indexTmp`、`cvtTmp` 三个区域；最小字节数按 PTO-ISA 文档公式计算；A5 不使用 tmp。 -- mask form 不使用 tmp，不应自动补 tmp。 -- A2/A3 level3 下 index / compare form 缺省 tmp 报错。 +- mask form 不使用 tmp。 +- A2/A3 index / compare form 必须显式提供 tmp；A5 index-form 不带 tmp。 TTRANS： From 1c2bcf5d8ddf9a00b565dec6e7bad4b6b8e73435 Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 6 Aug 2026 00:09:44 +0800 Subject: [PATCH 050/122] Fix tmp-aware sample bindings and rowexpand level3 --- ...oas-implicit-tmp-materialization-design.md | 2 +- .../Transforms/PTOMaterializeImplicitTmp.cpp | 8 +++-- .../pto/rowexpand_level3_no_tmp_preserved.pto | 29 +++++++++++++++++++ test/samples/Complex/mix_kernel.py | 2 +- test/samples/Rowmax/rowmax.py | 2 +- test/samples/Rowmin/rowmin.py | 2 +- test/samples/Rowprod/rowprod.py | 2 +- test/samples/Rowsum/rowsum.py | 2 +- test/samples/Sel/sel.py | 2 +- test/samples/Sels/sels.py | 4 +-- test/samples/Trans/trans.py | 2 +- test/samples/Xor/xor.py | 4 +-- test/samples/Xors/xors.py | 2 +- 13 files changed, 48 insertions(+), 15 deletions(-) create mode 100644 test/lit/pto/rowexpand_level3_no_tmp_preserved.pto diff --git a/docs/designs/ptoas-implicit-tmp-materialization-design.md b/docs/designs/ptoas-implicit-tmp-materialization-design.md index c9c3c9cdec..aaad9a2276 100644 --- a/docs/designs/ptoas-implicit-tmp-materialization-design.md +++ b/docs/designs/ptoas-implicit-tmp-materialization-design.md @@ -443,7 +443,7 @@ A5: - 如果已经有 `tmp`，pass 不修改，但 verifier 需要保证它只用于合法模式。 - A5 如果没有 `tmp`，pass 不修改。 - A2/A3 如果没有 `tmp`，且 op 是模式 1、当前 build level 会运行 memplan，则自动补 tmp。 -- A2/A3 如果没有 `tmp`，op 是模式 1、但当前 level3 会跳过 memplan，则报错，要求用户显式提供带地址的 tmp。 +- A2/A3 如果没有 `tmp`，op 是模式 1、但当前 level3 会跳过 memplan，则保留 no-tmp overload，由后端使用内部 8KB `TMP_UB_OFFSET`，避免 pass 生成无地址 tmp。 - 模式 2 不需要 tmp；pass 不应自动补 tmp，也不应强制改成带 tmp overload。 自动补 tmp 的 canonical shape 建议采用形状无关上界： diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp index d35c7147e3..138839984e 100644 --- a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -235,8 +235,12 @@ static LogicalResult materializeTRowExpandTmp(OpTy op, bool requireExplicitTmp, return success(); if (requireExplicitTmp) { - return op.emitOpError( - "requires explicit tmp for A2/A3 row-expand mode 1 when PlanMemory is skipped"); + // Row-expand is the one A2/A3 tmp-aware family whose no-tmp overload is + // still a valid backend contract: mode 1 falls back to pto-isa's internal + // 8KB TMP_UB_OFFSET scratch area, while mode 2 does not need tmp. Level3 + // inputs are already memory-planned by the frontend, so preserve their + // no-tmp form instead of creating an unaddressed alloc_tile. + return success(); } auto dstTy = dyn_cast(op.getDst().getType()); diff --git a/test/lit/pto/rowexpand_level3_no_tmp_preserved.pto b/test/lit/pto/rowexpand_level3_no_tmp_preserved.pto new file mode 100644 index 0000000000..ef1f6fa6ce --- /dev/null +++ b/test/lit/pto/rowexpand_level3_no_tmp_preserved.pto @@ -0,0 +1,29 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level3 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @rowexpand_level3_no_tmp_preserved() { + %addr0 = arith.constant 0 : i64 + %addr1 = arith.constant 4096 : i64 + %addr2 = arith.constant 8192 : i64 + %src0 = pto.alloc_tile addr = %addr0 : !pto.tile_buf + %src1 = pto.alloc_tile addr = %addr1 : !pto.tile_buf + %dst = pto.alloc_tile addr = %addr2 : !pto.tile_buf + pto.trowexpandsub ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + pto.trowexpanddiv ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK-LABEL: func.func @rowexpand_level3_no_tmp_preserved +// CHECK: pto.trowexpandsub ins(%{{.*}}, %{{.*}} : +// CHECK-NOT: pto.trowexpandsub ins(%{{.*}}, %{{.*}}, %{{.*}} : +// CHECK: pto.trowexpanddiv ins(%{{.*}}, %{{.*}} : +// CHECK-NOT: pto.trowexpanddiv ins(%{{.*}}, %{{.*}}, %{{.*}} : diff --git a/test/samples/Complex/mix_kernel.py b/test/samples/Complex/mix_kernel.py index a7f4e8f6ba..031d0f2a07 100644 --- a/test/samples/Complex/mix_kernel.py +++ b/test/samples/Complex/mix_kernel.py @@ -132,7 +132,7 @@ def build(M=32, N=32, K=32, TM=32, TN=32, TK=32): pto.TLoadOp(None, sv_out, ubTile) # pto.trowmax ins(%src, %tmp) outs(%dst) - pto.TRowMaxOp(ubTile, ubTmpTile, ubReduceTile) + pto.TRowMaxOp(ubTile, ubReduceTile, tmp=ubTmpTile) pto.TStoreOp(None, ubReduceTile, sv_reduce) func.ReturnOp([]) diff --git a/test/samples/Rowmax/rowmax.py b/test/samples/Rowmax/rowmax.py index 30819eca8a..a63f8b05b5 100644 --- a/test/samples/Rowmax/rowmax.py +++ b/test/samples/Rowmax/rowmax.py @@ -67,7 +67,7 @@ def build(): pto.TLoadOp(None, sv0, tb0) # pto.trowmax ins(%src, %tmp) outs(%dst) - pto.TRowMaxOp(tb0, tb_tmp, tb1) + pto.TRowMaxOp(tb0, tb1, tmp=tb_tmp) # %8 = subview on output tensor_view sv1 = pto.PartitionViewOp(tile_view_32x1, tv1, offsets=[c0, c0], sizes=[c32, c1]).result diff --git a/test/samples/Rowmin/rowmin.py b/test/samples/Rowmin/rowmin.py index a51efdf81c..5840666982 100644 --- a/test/samples/Rowmin/rowmin.py +++ b/test/samples/Rowmin/rowmin.py @@ -60,7 +60,7 @@ def build(): tb1 = pto.AllocTileOp(tile_buf_32x1).result pto.TLoadOp(None, sv0, tb0) - pto.TRowMinOp(tb0, tb_tmp, tb1) + pto.TRowMinOp(tb0, tb1, tmp=tb_tmp) sv1 = pto.PartitionViewOp(tile_view_32x1, tv1, offsets=[c0, c0], sizes=[c32, c1]).result pto.TStoreOp(None, tb1, sv1) diff --git a/test/samples/Rowprod/rowprod.py b/test/samples/Rowprod/rowprod.py index 15e7c55734..c10ebe64e6 100644 --- a/test/samples/Rowprod/rowprod.py +++ b/test/samples/Rowprod/rowprod.py @@ -57,7 +57,7 @@ def build(): tb1 = pto.AllocTileOp(tile_buf_32x1).result pto.TLoadOp(None, sv0, tb0) - pto.TRowProdOp(tb0, tb_tmp, tb1) + pto.TRowProdOp(tb0, tb1, tmp=tb_tmp) sv1 = pto.PartitionViewOp(tile_view_32x1, tv1, offsets=[c0, c0], sizes=[c32, c1]).result pto.TStoreOp(None, tb1, sv1) diff --git a/test/samples/Rowsum/rowsum.py b/test/samples/Rowsum/rowsum.py index d7c5842f82..4a0e10b724 100644 --- a/test/samples/Rowsum/rowsum.py +++ b/test/samples/Rowsum/rowsum.py @@ -67,7 +67,7 @@ def build(): pto.TLoadOp(None, sv0, tb0) # result=None, valid_dims=[] # pto.trowsum ins(%src, %tmp) outs(%dst) - pto.TRowSumOp(tb0, tb_tmp, tb1) + pto.TRowSumOp(tb0, tb1, tmp=tb_tmp) # %8 = subview on output tensor_view sv1 = pto.PartitionViewOp(tile_view_32x1, tv1, offsets=[c0, c0], sizes=[c32, c1]).result diff --git a/test/samples/Sel/sel.py b/test/samples/Sel/sel.py index f4ae49b5ac..2d2a04dcdb 100644 --- a/test/samples/Sel/sel.py +++ b/test/samples/Sel/sel.py @@ -77,7 +77,7 @@ def build(): pto.TLoadOp(None, sv2, tb2) # result=None # pto.tsel ins(%mask,%src0,%src1,%tmp) outs(%dst) - pto.TSelOp(tb0, tb1, tb2, tb_tmp, tb3) + pto.TSelOp(tb0, tb1, tb2, tb3, tmp=tb_tmp) # %8 = subview on output tensor_view sv3 = pto.PartitionViewOp(tile_view_f32, tv3, offsets=[c0, c0], sizes=[c32, c32]).result diff --git a/test/samples/Sels/sels.py b/test/samples/Sels/sels.py index 1e10484e6b..f1e10bf5f5 100644 --- a/test/samples/Sels/sels.py +++ b/test/samples/Sels/sels.py @@ -65,8 +65,8 @@ def build(): pto.TLoadOp(None, sv0, tb0) # result=None pto.TLoadOp(None, sv1, tb1) # result=None - # TSELS(mask=tb0, src=tb1, tmp=tb2, scalar=c64) - pto.TSelSOp(tb0, tb1, tb2, c64, tb2) + # TSELS(mask=tb0, src=tb1, dst=tb2, tmp=tb2, scalar=c64) + pto.TSelSOp(tb0, tb1, c64, tb2, tmp=tb2) # %8 = subview on output tensor_view sv2 = pto.PartitionViewOp(tile_view_32, tv2, offsets=[c0, c0], sizes=[c32, c32]).result diff --git a/test/samples/Trans/trans.py b/test/samples/Trans/trans.py index 9b17b31b1e..e948b2cd5c 100644 --- a/test/samples/Trans/trans.py +++ b/test/samples/Trans/trans.py @@ -63,7 +63,7 @@ def build(): pto.TLoadOp(None, sv0, tb_src) # transpose: ttrans ins(%src, %tmp) outs(%dst) - pto.TTransOp(tb_src, tb_tmp, tb_dst) + pto.TTransOp(tb_src, tb_dst, tmp=tb_tmp) # output subview sv1 = pto.PartitionViewOp(tile_view_32, tv1, offsets=[c0, c0], sizes=[c32, c32]).result diff --git a/test/samples/Xor/xor.py b/test/samples/Xor/xor.py index 86acc61a46..2a8dd54b83 100644 --- a/test/samples/Xor/xor.py +++ b/test/samples/Xor/xor.py @@ -64,8 +64,8 @@ def build(): pto.TLoadOp(None, sv_src0, tb_src0) # result=None pto.TLoadOp(None, sv_src1, tb_src1) - pto.TXorOp(tb_src0, tb_src1, tb_tmp, tb_dst) - pto.TXorOp(tb_src0, tb_src1, tb_dst, tb_dst) + pto.TXorOp(tb_src0, tb_src1, tb_dst, tmp=tb_tmp) + pto.TXorOp(tb_src0, tb_src1, tb_dst, tmp=tb_dst) # output subview sv_dst = pto.PartitionViewOp(tile_view_32, tv_dst, offsets=[c0, c0], sizes=[c32, c32]).result diff --git a/test/samples/Xors/xors.py b/test/samples/Xors/xors.py index 3ef70a27ea..43bb176a3c 100644 --- a/test/samples/Xors/xors.py +++ b/test/samples/Xors/xors.py @@ -61,7 +61,7 @@ def build(): pto.TLoadOp(None, sv_src, tb_src) # result=None - pto.TXorSOp(tb_src, scale, tb_tmp, tb_dst) + pto.TXorSOp(tb_src, scale, tb_dst, tmp=tb_tmp) # output subview sv_dst = pto.PartitionViewOp(tile_view_32, tv_dst, offsets=[c0, c0], sizes=[c32, c32]).result From e2622b2db8cd20bad483499cb2dac299835e037d Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 6 Aug 2026 08:52:59 +0800 Subject: [PATCH 051/122] Update rowexpand level3 tmp lit expectation --- ...lid.pto => trowexpandadd_level3_no_tmp_preserved.pto} | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) rename test/lit/pto/{trowexpand_implicit_tmp_level3_invalid.pto => trowexpandadd_level3_no_tmp_preserved.pto} (84%) diff --git a/test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto b/test/lit/pto/trowexpandadd_level3_no_tmp_preserved.pto similarity index 84% rename from test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto rename to test/lit/pto/trowexpandadd_level3_no_tmp_preserved.pto index 34f9d32a7c..5c53792ae5 100644 --- a/test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto +++ b/test/lit/pto/trowexpandadd_level3_no_tmp_preserved.pto @@ -6,19 +6,22 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: not ptoas --pto-arch=a3 --pto-level=level3 %s 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a3 --pto-level=level3 --emit-pto-ir %s 2>&1 | FileCheck %s module { - func.func @trowexpand_mode1_level3_requires_tmp() { + func.func @trowexpand_mode1_level3_preserves_no_tmp() { %addr0 = arith.constant 0 : i64 %addr1 = arith.constant 1024 : i64 %addr2 = arith.constant 2048 : i64 %src0 = pto.alloc_tile addr = %addr0 : !pto.tile_buf %src1 = pto.alloc_tile addr = %addr1 : !pto.tile_buf %dst = pto.alloc_tile addr = %addr2 : !pto.tile_buf - // CHECK: error: 'pto.trowexpandadd' op requires explicit tmp for A2/A3 row-expand mode 1 when PlanMemory is skipped pto.trowexpandadd ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) return } } + +// CHECK-LABEL: func.func @trowexpand_mode1_level3_preserves_no_tmp +// CHECK: pto.trowexpandadd ins(%{{.*}}, %{{.*}} : +// CHECK-NOT: pto.trowexpandadd ins(%{{.*}}, %{{.*}}, %{{.*}} : From 64eeb54cf4f41f9a365b66660d6278d9e98da297 Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 6 Aug 2026 10:19:14 +0800 Subject: [PATCH 052/122] Fix tmrgsort smoke tmp tile size --- .../npu/a5/src/st/smoke/testcase/tmrgsort/tmrgsort.pto | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tmrgsort/tmrgsort.pto b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tmrgsort/tmrgsort.pto index ed90162b9c..2131158a94 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tmrgsort/tmrgsort.pto +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tmrgsort/tmrgsort.pto @@ -107,7 +107,7 @@ func.func @TMRGSORT_f16_topk_1280_512(%src_ptr: !pto.ptr, %dst_ptr: !pto.pt %block1_tile = pto.alloc_tile : !pto.tile_buf %merge_tmp_tile = pto.alloc_tile - : !pto.tile_buf + : !pto.tile_buf %merge_dst_tile = pto.alloc_tile : !pto.tile_buf %ex_vec = arith.constant dense<0> : vector<4xi16> @@ -159,7 +159,7 @@ func.func @TMRGSORT_f16_topk_1280_512(%src_ptr: !pto.ptr, %dst_ptr: !pto.pt pto.tmrgsort ins(%block0_tile, %block1_tile, %merge_tmp_tile {exhausted = false} : !pto.tile_buf, !pto.tile_buf, - !pto.tile_buf) + !pto.tile_buf) outs(%merge_dst_tile, %ex_vec : !pto.tile_buf, vector<4xi16>) From 81bdb358967b379c602dc64bda7036da807ccb23 Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 6 Aug 2026 11:32:19 +0800 Subject: [PATCH 053/122] Fix PTODSL ttrans tmp builder call --- ptodsl/ptodsl/_ops.py | 4 ++-- ptodsl/tests/test_vector_cube_ops.py | 11 +++++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 5ec91c9a66..bf4110d623 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -3136,10 +3136,10 @@ def tmov(src, dst, *, mode=None): def ttrans(src, tmp, dst): """``pto.ttrans ins(src, tmp) outs(dst)`` – tile transpose (DPS).""" - _pto.ttrans( + _pto.TTransOp( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index ab134cb99e..acc4095989 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -764,6 +764,17 @@ def test_tile_row_reductions_expose_optional_tmp_and_synthesize_one(self): getattr(pto.tile, name)(src, dst, tmp=tmp) low_level_op.assert_called_once_with(src, tmp, dst) + def test_tile_transpose_wrapper_uses_tmp_keyword_builder(self): + src = object() + tmp = object() + dst = object() + + with patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops._pto, "TTransOp") as ttrans_op: + pto.tile.transpose(src, tmp, dst) + + ttrans_op.assert_called_once_with(src, dst, tmp=tmp) + def test_tile_sort_gather_wrappers_call_low_level_ops(self): src = object() idx = object() From 657a4ba996aa5d63a13646c84bfc52ed728bb8a1 Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 6 Aug 2026 23:23:44 +0800 Subject: [PATCH 054/122] Fix implicit tmp contracts --- lib/PTO/IR/PTO.cpp | 180 +++++++++++- .../Transforms/PTOMaterializeImplicitTmp.cpp | 267 ++++++++---------- test/lit/pto/implicit_tmp_a5_skip_no_tmp.pto | 52 ++++ test/lit/pto/implicit_tmp_arg_reductions.pto | 19 +- test/lit/pto/implicit_tmp_row_reductions.pto | 19 +- .../pto/implicit_tmp_xor_materialization.pto | 13 +- .../pto/plan_memory_inplace_forbid_alias.pto | 4 +- test/lit/pto/select_tile_native.pto | 4 +- test/lit/pto/tsel_bf16.pto | 4 +- test/lit/pto/tsel_tmp_contract_a3_invalid.pto | 24 ++ 10 files changed, 398 insertions(+), 188 deletions(-) create mode 100644 test/lit/pto/implicit_tmp_a5_skip_no_tmp.pto create mode 100644 test/lit/pto/tsel_tmp_contract_a3_invalid.pto diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 202fefe1e3..7c1f269e72 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -122,6 +122,9 @@ static bool isKnownZeroOrUnitExtent(int64_t value); static bool isByteIntegerType(Type ty); static LogicalResult verifyTileBufCommon(Operation *op, Type ty, StringRef name, bool allowLowPrecision = false); +static LogicalResult verifyTmpCapacityAtLeast(Operation *op, Type tmpTy, + uint64_t requiredBytes, + StringRef tmpName = "tmp"); namespace { struct PTOInlinerInterface : public DialectInlinerInterface { @@ -2144,11 +2147,17 @@ static LogicalResult verifyTRowReductionWithTmpCommon(Operation *op, Type srcTy, return failure(); if (getElemTy(srcTy) != getElemTy(dstTy)) return op->emitOpError("expects src and dst to have the same element type"); + if (getTargetArch(op) != PTOArch::A5 && + getElemTy(srcTy) != getElemTy(tmpTy)) + return op->emitOpError("expects A2/A3 tmp to have the same element type as src and dst"); if (failed(verifyRowReductionValidRegion(op, srcTy, dstTy, /*allowEmptyMarker=*/true))) return failure(); if (!isSupportedRowReductionElemType(getElemTy(srcTy))) return op->emitOpError(elemTypeError); + if (getTargetArch(op) != PTOArch::A5 && + failed(verifyTmpCapacityAtLeast(op, tmpTy, 32))) + return failure(); return success(); } @@ -2197,7 +2206,7 @@ static LogicalResult verifyTColArgTmpA2A3(Operation *op, Type srcTy, return failure(); if (hasExactKnownValidShape(srcTy, tmpTy)) - return success(); + return verifyTmpCapacityAtLeast(op, tmpTy, 32); auto srcValid = getValidShapeVec(srcTy); auto tmpValid = getValidShapeVec(tmpTy); @@ -2214,7 +2223,7 @@ static LogicalResult verifyTColArgTmpA2A3(Operation *op, Type srcTy, << "expects A2/A3 tmp valid_shape[1] to be at least " << *minStride << " for src valid_shape[1] = " << srcValid[1]; } - return success(); + return verifyTmpCapacityAtLeast(op, tmpTy, 32); } static LogicalResult verifyTColArgReductionOpA2A3(Operation *op, Type srcTy, @@ -2307,7 +2316,7 @@ static LogicalResult verifyTRowArgTmpA2A3(Operation *op, Type srcTy, return failure(); if (hasExactKnownValidShape(srcTy, tmpTy)) - return success(); + return verifyTmpCapacityAtLeast(op, tmpTy, 32); auto srcShape = getShapeVec(srcTy); auto tmpShape = getShapeVec(tmpTy); @@ -2335,7 +2344,7 @@ static LogicalResult verifyTRowArgTmpA2A3(Operation *op, Type srcTy, return op->emitOpError() << "expects A2/A3 tmp DN layout to have valid_shape[0] >= " << (srcValid[0] * 2); - return success(); + return verifyTmpCapacityAtLeast(op, tmpTy, 32); } if (!layout || *layout != pto::Layout::ND) @@ -2349,7 +2358,7 @@ static LogicalResult verifyTRowArgTmpA2A3(Operation *op, Type srcTy, if (tmpValid[1] != ShapedType::kDynamic && tmpValid[1] < 2) return op->emitOpError( "expects A2/A3 tmp valid_shape[1] to be at least 2 in the small-col ND path"); - return success(); + return verifyTmpCapacityAtLeast(op, tmpTy, 32); } if (failed(verifyVecTileCommon(op, tmpTy, "tmp"))) @@ -2369,7 +2378,7 @@ static LogicalResult verifyTRowArgTmpA2A3(Operation *op, Type srcTy, << "expects A2/A3 tmp valid_shape[1] to be at least " << *minStride << " for src valid_shape[1] = " << srcValid[1]; } - return success(); + return verifyTmpCapacityAtLeast(op, tmpTy, 32); } static LogicalResult verifyTRowArgReductionOpA2A3(Operation *op, Type srcTy, @@ -4164,6 +4173,20 @@ static std::optional getStaticByteSize(Type ty) { return total; } +static LogicalResult verifyTmpCapacityAtLeast(Operation *op, Type tmpTy, + uint64_t requiredBytes, + StringRef tmpName) { + auto actualBytes = getStaticByteSize(tmpTy); + if (!actualBytes) + return op->emitOpError() + << "expects " << tmpName << " to have statically known byte capacity"; + if (*actualBytes < requiredBytes) + return op->emitOpError() + << "expects " << tmpName << " capacity to be at least " + << requiredBytes << " bytes, but got " << *actualBytes << " bytes"; + return success(); +} + static std::optional getPTOMemorySpaceEnum(Type ty) { if (auto ptr = dyn_cast(ty)) return ptr.getMemorySpace().getAddressSpace(); @@ -6376,6 +6399,19 @@ LogicalResult pto::TColSumOp::verify() { return emitOpError("expects src/tmp/dst element types to match"); if (failed(verifyTColSumTmpStride(*this, srcTy, tmpTy, getIsBinary()))) return failure(); + if (getIsBinary()) { + auto srcValid = getValidShapeVec(srcTy); + auto elemBytes = getElemByteSize(getElemTy(srcTy)); + if (srcValid.size() != 2 || srcValid[0] == ShapedType::kDynamic || + srcValid[1] == ShapedType::kDynamic || elemBytes == 0) + return emitOpError( + "expects static src valid_shape and element size to verify tcolsum tmp"); + uint64_t requiredBytes = + static_cast(ceilDivInt64(srcValid[0], 2)) * + static_cast(srcValid[1]) * elemBytes; + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, requiredBytes))) + return failure(); + } } if (getElemTy(srcTy) != getElemTy(dstTy)) return emitOpError("expects src/dst element types to match"); @@ -6405,6 +6441,19 @@ LogicalResult pto::TColSumOp::verify() { return emitOpError("expects src/tmp/dst element types to match"); if (failed(verifyTColSumTmpStride(*this, srcTy, tmpTy, getIsBinary()))) return failure(); + if (getIsBinary()) { + auto srcValid = getValidShapeVec(srcTy); + auto elemBytes = getElemByteSize(getElemTy(srcTy)); + if (srcValid.size() != 2 || srcValid[0] == ShapedType::kDynamic || + srcValid[1] == ShapedType::kDynamic || elemBytes == 0) + return emitOpError( + "expects static src valid_shape and element size to verify tcolsum tmp"); + uint64_t requiredBytes = + static_cast(ceilDivInt64(srcValid[0], 2)) * + static_cast(srcValid[1]) * elemBytes; + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, requiredBytes))) + return failure(); + } } if (getElemTy(srcTy) != getElemTy(dstTy)) return emitOpError("expects src/dst element types to match"); @@ -10473,6 +10522,18 @@ mlir::LogicalResult mlir::pto::TPReluOp::verify() { << "expects A2/A3 tmp valid_shape[1] to be at least ceil(dst valid_shape[1] / 8) (" << packedMaskCols << ")"; } + if (dstValid[0] == ShapedType::kDynamic || + dstValid[1] == ShapedType::kDynamic) + return emitOpError( + "expects A2/A3 tprelu dst valid_shape to be static when tmp is provided"); + int64_t packedCols = std::max( + 32, llvm::divideCeil(llvm::divideCeil(dstValid[1], int64_t{8}), + int64_t{32}) * + 32); + if (failed(verifyTmpCapacityAtLeast( + *this, tt, static_cast(dstValid[0] + 1) * + static_cast(packedCols)))) + return failure(); if (auto arch = getVerifierArchName(getOperation()); arch && arch->equals_insensitive("a3")) { if (getSrc0() == getSrc1() || getSrc0() == getTmp() || getSrc0() == getDst() || @@ -10740,7 +10801,11 @@ mlir::LogicalResult mlir::pto::TQuantOp::verify() { return emitOpError() << "expects A2/A3 tmp to have the same shape as src"; if (failed(verifyTileBufSameValidShape(*this, srcTy, tmpTy, "src", "tmp"))) return failure(); - return success(); + auto requiredBytes = getStaticByteSize(srcTy); + if (!requiredBytes) + return emitOpError( + "expects A2/A3 tquant src shape to be static when tmp is provided"); + return verifyTmpCapacityAtLeast(*this, tmpTy, *requiredBytes); }; if (getTmp() && failed(verifyA2A3Tmp(getTmp().getType()))) return failure(); @@ -11040,6 +11105,16 @@ mlir::LogicalResult mlir::pto::TRemOp::verify() { if (dstValid[1] != ShapedType::kDynamic && tmpValid[1] != ShapedType::kDynamic && tmpValid[1] < dstValid[1]) return emitOpError("expects A2/A3 tmp valid columns to cover dst valid columns"); + auto dstShape = getShapeVec(dstTy); + auto elemBytes = getElemByteSize(elem); + if (dstShape.size() != 2 || dstShape[1] == ShapedType::kDynamic || + elemBytes == 0) + return emitOpError( + "expects A2/A3 trem dst shape and element size to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast( + *this, tmpTy, static_cast(2) * + static_cast(dstShape[1]) * elemBytes))) + return failure(); if (!(elem.isInteger(32) || elem.isF32())) return emitOpError("expects A2/A3 trem element type to be i32/f32"); return success(); @@ -11110,6 +11185,15 @@ mlir::LogicalResult mlir::pto::TRemSOp::verify() { if (dstValid[1] != ShapedType::kDynamic && tmpValid[1] != ShapedType::kDynamic && tmpValid[1] < dstValid[1]) return emitOpError("expects A2/A3 tmp valid columns to cover dst valid columns"); + auto dstShape = getShapeVec(td); + auto elemBytes = getElemByteSize(elem); + if (dstShape.size() != 2 || dstShape[1] == ShapedType::kDynamic || + elemBytes == 0) + return emitOpError( + "expects A2/A3 trems dst shape and element size to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast( + *this, tt, static_cast(dstShape[1]) * elemBytes))) + return failure(); if (!(elem.isInteger(32) || elem.isF32())) return emitOpError("expects A2/A3 trems element type to be i32/f32"); return success(); @@ -11218,6 +11302,14 @@ mlir::LogicalResult mlir::pto::TPowOp::verify() { return failure(); if (failed(verifyTPowTmpShape(getOperation(), tmpTy, dstTy))) return failure(); + if (getTargetArch(getOperation()) != PTOArch::A5) { + auto requiredBytes = getStaticByteSize(dstTy); + if (!requiredBytes) + return emitOpError( + "expects A2/A3 tpow dst shape to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, *requiredBytes))) + return failure(); + } } return success(); } @@ -11276,6 +11368,14 @@ mlir::LogicalResult mlir::pto::TPowSOp::verify() { return failure(); if (failed(verifyTPowTmpShape(getOperation(), tmpTy, dstTy))) return failure(); + if (getTargetArch(getOperation()) != PTOArch::A5) { + auto requiredBytes = getStaticByteSize(dstTy); + if (!requiredBytes) + return emitOpError( + "expects A2/A3 tpows dst shape to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, *requiredBytes))) + return failure(); + } } return success(); } @@ -12934,6 +13034,9 @@ mlir::LogicalResult mlir::pto::TSelOp::verify() { failed(verifyTileBufCommon(*this, t1, "src1")) || failed(verifyTileBufCommon(*this, td, "dst"))) return failure(); + if (getTmp() && + failed(verifyVecTileCommon(*this, getTmp().getType(), "tmp"))) + return failure(); Type srcElem = getElemTy(t0); Type src1Elem = getElemTy(t1); @@ -12967,6 +13070,18 @@ mlir::LogicalResult mlir::pto::TSelOp::verify() { if (!ok) return emitOpError( "expects A2/A3 tsel src0, src1, and dst element type to be i16/i32/f16/bf16/f32"); + if (getTmp()) { + Type tmpTy = getTmp().getType(); + auto tmpElem = dyn_cast(getElemTy(tmpTy)); + if (!tmpElem || tmpElem.getWidth() != 32) + return emitOpError("expects A2/A3 tsel tmp element type to be i32"); + unsigned elemBits = getPTOStorageElemBitWidth(elem); + if (elemBits != 16 && elemBits != 32) + return emitOpError("expects A2/A3 tsel data element type to be 16 or 32 bits"); + uint64_t minBytes = elemBits == 16 ? 64 : 32; + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, minBytes))) + return failure(); + } return success(); }; @@ -13025,6 +13140,22 @@ mlir::LogicalResult mlir::pto::TSelSOp::verify() { if (!isRowMajorTileBuf(tSrc) || !isRowMajorTileBuf(tDst)) return emitOpError("expects src and dst to use row-major layout"); Type elem = *elemOr; + if (getTmp()) { + Type tmpTy = getTmp().getType(); + if (getElemTy(tmpTy) != elem) + return emitOpError("expects A2/A3 tsels tmp to have the same element type as src and dst"); + if (!isRowMajorTileBuf(tmpTy)) + return emitOpError("expects A2/A3 tsels tmp to use row-major layout"); + auto srcShape = getShapeVec(tSrc); + if (srcShape.size() != 2 || srcShape[1] == ShapedType::kDynamic) + return emitOpError( + "expects A2/A3 tsels src shape to be static when tmp is provided"); + auto elemBytes = getElemByteSize(elem); + if (elemBytes == 0 || + failed(verifyTmpCapacityAtLeast( + *this, tmpTy, static_cast(srcShape[1]) * elemBytes))) + return failure(); + } bool ok = elem.isF16() || elem.isF32(); if (auto it = mlir::dyn_cast(elem)) ok = (it.getWidth() == 16 || it.getWidth() == 32); @@ -13104,6 +13235,15 @@ mlir::LogicalResult mlir::pto::TSort32Op::verify() { if (getTmp() && failed(verifyVecTileCommon(*this, getTmp().getType(), "tmp"))) return failure(); + if (getTmp() && getTargetArch(getOperation()) != PTOArch::A5) { + auto requiredBytes = getStaticByteSize(srcTy); + if (!requiredBytes) + return emitOpError( + "expects A2/A3 tsort32 src shape to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast(*this, getTmp().getType(), + *requiredBytes))) + return failure(); + } auto srcElem = getElemTy(srcTy); auto dstElem = getElemTy(dstTy); @@ -13329,6 +13469,18 @@ mlir::LogicalResult mlir::pto::TTransOp::verify() { }; if (!isAllowedWidthType(srcElem)) return emitOpError() << "expects transpose element type to match the supported set for its width"; + if (tmpTy) { + uint64_t requiredBytes = 32; + if (ttransUsesTmp(srcTy, dstTy)) { + auto srcBytes = getStaticByteSize(srcTy); + if (!srcBytes) + return emitOpError( + "expects A2/A3 transpose src shape to be static when tmp is used"); + requiredBytes = *srcBytes; + } + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, requiredBytes))) + return failure(); + } return mlir::success(); }; auto verifyA5 = [&]() -> LogicalResult { @@ -13359,6 +13511,8 @@ mlir::LogicalResult mlir::pto::TTransOp::verify() { }; if (!isAllowedWidthType(srcElem)) return emitOpError() << "expects transpose element type to match the supported set for its width"; + if (tmpTy && failed(verifyTmpCapacityAtLeast(*this, tmpTy, 32))) + return failure(); auto checkAlignedMajor = [&](Type ty, StringRef name) -> LogicalResult { auto tb = mlir::dyn_cast(ty); if (!tb) @@ -13497,6 +13651,12 @@ mlir::LogicalResult mlir::pto::TXorOp::verify() { if (failed(verifyTileBufSameValidShape( *this, tmpTy, getDst().getType(), "tmp", "dst"))) return failure(); + auto requiredBytes = getStaticByteSize(getDst().getType()); + if (!requiredBytes) + return emitOpError( + "expects A2/A3 txor dst shape to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, *requiredBytes))) + return failure(); } auto it = mlir::dyn_cast(elem); if (!it || (it.getWidth() != 8 && it.getWidth() != 16 && @@ -13542,6 +13702,12 @@ mlir::LogicalResult mlir::pto::TXorSOp::verify() { "expects tmp to have the same element type as src and dst"); if (!isRowMajorTileBuf(tmpTy)) return emitOpError("expects tmp to use row-major layout"); + auto requiredBytes = getStaticByteSize(getDst().getType()); + if (!requiredBytes) + return emitOpError( + "expects A2/A3 txors dst shape to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, *requiredBytes))) + return failure(); } auto it = mlir::dyn_cast(elem); if (!it || (it.getWidth() != 8 && it.getWidth() != 16)) diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp index 138839984e..b832503501 100644 --- a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -19,6 +19,8 @@ #include "mlir/Pass/Pass.h" #include "llvm/ADT/TypeSwitch.h" +#include + using namespace mlir; namespace { @@ -125,6 +127,21 @@ static void copyAttrsExceptOperandSegments(Operation *from, OperationState &to) } } +static void rebuildWithOperands(Operation *op, ArrayRef operands, + std::optional> segments) { + OpBuilder builder(op); + OperationState state(op->getLoc(), op->getName()); + state.addOperands(operands); + if (op->hasTrait()) { + assert(segments && "AttrSizedOperandSegments op must supply segments"); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr(*segments)); + } + copyAttrsExceptOperandSegments(op, state); + builder.create(state); + op->erase(); +} + static bool validShapesCompatible(ArrayRef lhs, ArrayRef rhs) { if (lhs.size() != rhs.size()) @@ -205,20 +222,23 @@ static pto::TileBufType makeTRowExpandTmpType(MLIRContext *ctx, makeRowMajorNoneBoxConfig(ctx)); } +static FailureOr makeA5PlaceholderTmpType( + MLIRContext *ctx, Value like, Type elementType = {}) { + auto likeTy = dyn_cast(like.getType()); + if (!likeTy) + return failure(); + if (!elementType) + elementType = likeTy.getElementType(); + auto elemBytes = getElemBytes(elementType); + if (!elemBytes || *elemBytes <= 0) + return failure(); + int64_t cols = std::max(1, 32 / *elemBytes); + return makeVecTmpType(ctx, {1, cols}, elementType, {1, cols}); +} + static void replaceTRowExpandBinaryOpWithTmp(Operation *op, Value src0, Value src1, Value tmp, Value dst) { - OpBuilder builder(op); - OperationState state(op->getLoc(), op->getName()); - state.addOperands({src0, src1, tmp, dst}); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr({1, 1, 1, 1})); - for (NamedAttribute attr : op->getAttrs()) { - if (attr.getName() == "operandSegmentSizes") - continue; - state.addAttribute(attr.getName(), attr.getValue()); - } - builder.create(state); - op->erase(); + rebuildWithOperands(op, {src0, src1, tmp, dst}, ArrayRef{1, 1, 1, 1}); } template @@ -283,11 +303,8 @@ static LogicalResult replaceTColSumWithTmp(pto::TColSumOp op, if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getSrc(), *tmp, op.getDst()}); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), {op.getSrc(), *tmp, op.getDst()}, + ArrayRef{1, 1, 1}); return success(); } @@ -309,19 +326,13 @@ static LogicalResult replaceTQuantWithTmp(pto::TQuantOp op, if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); SmallVector operands{op.getSrc(), op.getFp()}; if (op.getOffset()) operands.push_back(op.getOffset()); operands.push_back(*tmp); operands.push_back(op.getDst()); - state.addOperands(operands); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr( - {1, 1, op.getOffset() ? 1 : 0, 1, 1})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), operands, + ArrayRef{1, 1, op.getOffset() ? 1 : 0, 1, 1}); return success(); } @@ -347,11 +358,9 @@ static LogicalResult replaceTPowWithTmp(pto::TPowOp op, if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getBase(), op.getExp(), op.getDst(), *tmp}); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), + {op.getBase(), op.getExp(), op.getDst(), *tmp}, + ArrayRef{1, 1, 1, 1}); return success(); } @@ -372,68 +381,9 @@ static LogicalResult replaceTPowSWithTmp(pto::TPowSOp op, if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getSrc(), op.getScalar(), op.getDst(), *tmp}); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); - return success(); -} - -[[maybe_unused]] static LogicalResult -replaceTGatherWithTmp(pto::TGatherOp op, bool requireExplicitTmp, - MLIRContext *ctx) { - if (op.getTmp() || op.hasMaskForm()) - return success(); - if (!op.hasIndexForm() && !op.hasCompareForm()) - return success(); - if (requireExplicitTmp) - return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); - - FailureOr tmpType = failure(); - if (op.hasIndexForm()) { - tmpType = makeSameShapeTmpType(ctx, op.getIndices()); - } else { - auto srcTy = dyn_cast(op.getSrc().getType()); - auto dstTy = dyn_cast(op.getDst().getType()); - if (!srcTy || !dstTy) - return op.emitOpError( - "expects tile_buf operands when materializing compare-form tgather tmp"); - auto srcShape = getShapeVec(op.getSrc().getType()); - if (srcShape.size() != 2 || hasDynamicDim(srcShape)) - return op.emitOpError( - "requires static src shape to materialize compare-form tgather tmp"); - int64_t bytes = srcShape[0] * srcShape[1] * 4 + srcShape[0] * 4; - tmpType = makeVecTmpType(ctx, {1, bytes}, IntegerType::get(ctx, 8), - {1, bytes}); - } - if (failed(tmpType)) - return op.emitOpError( - "requires static tile_buf indices/src to materialize implicit tgather tmp"); - - OpBuilder builder(op); - FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); - if (failed(tmp)) - return failure(); - - OperationState state(op.getLoc(), op->getName()); - SmallVector operands{op.getSrc(), op.getDst()}; - if (op.getCdst()) - operands.push_back(op.getCdst()); - if (op.getIndices()) - operands.push_back(op.getIndices()); - operands.push_back(*tmp); - if (op.getKValue()) - operands.push_back(op.getKValue()); - state.addOperands(operands); - state.addAttribute( - "operandSegmentSizes", - builder.getDenseI32ArrayAttr({1, 1, op.getCdst() ? 1 : 0, - op.getIndices() ? 1 : 0, 1, - op.getKValue() ? 1 : 0})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), + {op.getSrc(), op.getScalar(), op.getDst(), *tmp}, + ArrayRef{1, 1, 1, 1}); return success(); } @@ -459,13 +409,8 @@ static LogicalResult replaceTSort32WithTmp(pto::TSort32Op op, if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getSrc(), op.getIdx(), *tmp, op.getDst()}); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr({1, 1, 1, 1})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), {op.getSrc(), op.getIdx(), *tmp, op.getDst()}, + ArrayRef{1, 1, 1, 1}); return success(); } @@ -475,10 +420,14 @@ static LogicalResult replaceRowReductionWithTmp(OpTy op, MLIRContext *ctx) { if (op.getTmp()) return success(); - if (requireExplicitTmp) + + bool isA5 = pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5; + if (requireExplicitTmp && !isA5) return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); - FailureOr tmpType = makeSameShapeTmpType(ctx, op.getSrc()); + FailureOr tmpType = + isA5 ? makeA5PlaceholderTmpType(ctx, op.getSrc()) + : makeSameShapeTmpType(ctx, op.getSrc()); if (failed(tmpType)) return op.emitOpError( "requires static tile_buf src to materialize implicit row-reduction tmp"); @@ -487,13 +436,8 @@ static LogicalResult replaceRowReductionWithTmp(OpTy op, if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getSrc(), *tmp, op.getDst()}); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr({1, 1, 1})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), {op.getSrc(), *tmp, op.getDst()}, + ArrayRef{1, 1, 1}); return success(); } @@ -502,9 +446,12 @@ static LogicalResult replaceTXorWithTmp(pto::TXorOp op, MLIRContext *ctx) { if (op.getTmp()) return success(); - if (requireExplicitTmp) + bool isA5 = pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5; + if (requireExplicitTmp && !isA5) return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); - FailureOr tmpType = makeSameShapeTmpType(ctx, op.getDst()); + FailureOr tmpType = + isA5 ? makeA5PlaceholderTmpType(ctx, op.getDst()) + : makeSameShapeTmpType(ctx, op.getDst()); if (failed(tmpType)) return op.emitOpError( "requires static tile_buf dst to materialize implicit txor tmp"); @@ -512,13 +459,9 @@ static LogicalResult replaceTXorWithTmp(pto::TXorOp op, FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getSrc0(), op.getSrc1(), *tmp, op.getDst()}); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr({1, 1, 1, 1})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), + {op.getSrc0(), op.getSrc1(), *tmp, op.getDst()}, + ArrayRef{1, 1, 1, 1}); return success(); } @@ -527,9 +470,12 @@ static LogicalResult replaceTXorSWithTmp(pto::TXorSOp op, MLIRContext *ctx) { if (op.getTmp()) return success(); - if (requireExplicitTmp) + bool isA5 = pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5; + if (requireExplicitTmp && !isA5) return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); - FailureOr tmpType = makeSameShapeTmpType(ctx, op.getDst()); + FailureOr tmpType = + isA5 ? makeA5PlaceholderTmpType(ctx, op.getDst()) + : makeSameShapeTmpType(ctx, op.getDst()); if (failed(tmpType)) return op.emitOpError( "requires static tile_buf dst to materialize implicit txors tmp"); @@ -537,13 +483,9 @@ static LogicalResult replaceTXorSWithTmp(pto::TXorSOp op, FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getSrc(), op.getScalar(), *tmp, op.getDst()}); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr({1, 1, 1, 1})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), + {op.getSrc(), op.getScalar(), *tmp, op.getDst()}, + ArrayRef{1, 1, 1, 1}); return success(); } @@ -566,13 +508,7 @@ static LogicalResult replaceFixedDpsOpWithTmp( else finalOperands.push_back(*tmp); } - OperationState state(op->getLoc(), op->getName()); - state.addOperands(finalOperands); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr(operandSegments)); - copyAttrsExceptOperandSegments(op, state); - builder.create(state); - op->erase(); + rebuildWithOperands(op, finalOperands, operandSegments); (void)opName; return success(); } @@ -610,60 +546,83 @@ static LogicalResult materializeFixedMandatoryTmp(Operation *op, .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) return success(); - auto type = makeTPReluTmpType(ctx, typedOp.getDst()); + bool isA5 = + pto::getTargetArch(op) == pto::PTOArch::A5; + auto type = isA5 ? makeA5PlaceholderTmpType( + ctx, typedOp.getDst(), + IntegerType::get(ctx, 8)) + : makeTPReluTmpType(ctx, typedOp.getDst()); if (failed(type)) return typedOp.emitOpError( "requires static tile_buf dst to materialize implicit tprelu tmp"); return replaceFixedDpsOpWithTmp( op, {typedOp.getSrc0(), typedOp.getSrc1(), Value(), typedOp.getDst()}, - *type, {1, 1, 1, 1}, requireExplicitTmp, "tprelu"); + *type, {1, 1, 1, 1}, isA5 ? false : requireExplicitTmp, + "tprelu"); }) .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) return success(); - auto type = makeRowsTmpType(ctx, typedOp.getDst(), 2); + bool isA5 = + pto::getTargetArch(op) == pto::PTOArch::A5; + auto type = isA5 ? makeA5PlaceholderTmpType(ctx, typedOp.getDst()) + : makeRowsTmpType(ctx, typedOp.getDst(), 2); if (failed(type)) return typedOp.emitOpError( "requires static tile_buf dst to materialize implicit trem tmp"); return replaceFixedDpsOpWithTmp( op, {typedOp.getSrc0(), typedOp.getSrc1(), Value(), typedOp.getDst()}, - *type, {1, 1, 1, 1}, requireExplicitTmp, "trem"); + *type, {1, 1, 1, 1}, isA5 ? false : requireExplicitTmp, "trem"); }) .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) return success(); - auto type = makeRowsTmpType(ctx, typedOp.getDst(), 1); + bool isA5 = + pto::getTargetArch(op) == pto::PTOArch::A5; + auto type = isA5 ? makeA5PlaceholderTmpType(ctx, typedOp.getDst()) + : makeRowsTmpType(ctx, typedOp.getDst(), 1); if (failed(type)) return typedOp.emitOpError( "requires static tile_buf dst to materialize implicit trems tmp"); return replaceFixedDpsOpWithTmp( op, {typedOp.getSrc(), typedOp.getScalar(), Value(), typedOp.getDst()}, - *type, {1, 1, 1, 1}, requireExplicitTmp, "trems"); + *type, {1, 1, 1, 1}, isA5 ? false : requireExplicitTmp, "trems"); }) .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) return success(); - auto type = makeVecTmpType(ctx, {1, 16}, IntegerType::get(ctx, 32), - {1, 16}); + bool isA5 = + pto::getTargetArch(op) == pto::PTOArch::A5; + auto type = isA5 ? makeA5PlaceholderTmpType( + ctx, typedOp.getDst(), + IntegerType::get(ctx, 32)) + : makeVecTmpType(ctx, {1, 16}, + IntegerType::get(ctx, 32), {1, 16}); + if (failed(type)) + return typedOp.emitOpError( + "requires static tile_buf dst to materialize implicit tsel tmp"); return replaceFixedDpsOpWithTmp( op, {typedOp.getMask(), typedOp.getSrc0(), typedOp.getSrc1(), Value(), typedOp.getDst()}, - type, {1, 1, 1, 1, 1}, requireExplicitTmp, "tsel"); + *type, {1, 1, 1, 1, 1}, isA5 ? false : requireExplicitTmp, "tsel"); }) .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) return success(); - auto type = makeRowsTmpType(ctx, typedOp.getSrc(), 1); + bool isA5 = + pto::getTargetArch(op) == pto::PTOArch::A5; + auto type = isA5 ? makeA5PlaceholderTmpType(ctx, typedOp.getSrc()) + : makeRowsTmpType(ctx, typedOp.getSrc(), 1); if (failed(type)) return typedOp.emitOpError( "requires static tile_buf src to materialize implicit tsels tmp"); return replaceFixedDpsOpWithTmp( op, {typedOp.getMask(), typedOp.getSrc(), Value(), typedOp.getScalar(), typedOp.getDst()}, - *type, {1, 1, 1, 1, 1}, requireExplicitTmp, "tsels"); + *type, {1, 1, 1, 1, 1}, isA5 ? false : requireExplicitTmp, "tsels"); }) .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) @@ -791,18 +750,14 @@ static LogicalResult materializeTMrgSortTmp(pto::TMrgSortOp op, FailureOr tmp = createAllocTmp(builder, op.getLoc(), tmpType); if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands(operands); - state.addOperands(op.getDsts()); - state.addOperands(*tmp); - state.addOperands(op.getExcuted()); - state.addAttribute( - "operandSegmentSizes", - builder.getDenseI32ArrayAttr( - {static_cast(op.getSrcs().size()), 0, 1, 1, 1})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + SmallVector finalOperands; + finalOperands.append(operands.begin(), operands.end()); + finalOperands.append(op.getDsts().begin(), op.getDsts().end()); + finalOperands.push_back(*tmp); + finalOperands.push_back(op.getExcuted()); + rebuildWithOperands(op.getOperation(), finalOperands, + ArrayRef{static_cast(op.getSrcs().size()), + 0, 1, 1, 1}); return success(); } diff --git a/test/lit/pto/implicit_tmp_a5_skip_no_tmp.pto b/test/lit/pto/implicit_tmp_a5_skip_no_tmp.pto new file mode 100644 index 0000000000..3b66956054 --- /dev/null +++ b/test/lit/pto/implicit_tmp_a5_skip_no_tmp.pto @@ -0,0 +1,52 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: sed -E 's/ addr = %[A-Za-z0-9_]+//g' %s > %t.level2.pto && ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %t.level2.pto 2>&1 | FileCheck %s --check-prefix=A5 +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 + +module { + func.func @a5_skip_no_tmp_overloads() attributes {pto.kernel_kind = #pto.kernel_kind} { + %c0_i32 = arith.constant 0 : i32 + %a0 = arith.constant 0 : i64 + %a256 = arith.constant 256 : i64 + %a512 = arith.constant 512 : i64 + %a768 = arith.constant 768 : i64 + %a1024 = arith.constant 1024 : i64 + %a1280 = arith.constant 1280 : i64 + %a1536 = arith.constant 1536 : i64 + %a1792 = arith.constant 1792 : i64 + %a2048 = arith.constant 2048 : i64 + %seq = pto.alloc_tile addr = %a0 : !pto.tile_buf + pto.tci ins(%c0_i32 : i32) + outs(%seq : !pto.tile_buf) + + %cvt_src = pto.alloc_tile addr = %a256 : !pto.tile_buf + %cvt_dst = pto.alloc_tile addr = %a512 : !pto.tile_buf + pto.tcvt ins(%cvt_src : !pto.tile_buf) + outs(%cvt_dst : !pto.tile_buf) + + %q_src = pto.alloc_tile addr = %a768 : !pto.tile_buf + %q_fp = pto.alloc_tile addr = %a1024 : !pto.tile_buf + %q_dst = pto.alloc_tile addr = %a1280 : !pto.tile_buf + pto.tquant ins(%q_src, %q_fp : !pto.tile_buf, !pto.tile_buf) + outs(%q_dst : !pto.tile_buf) {quant_type = #pto} + + %re_src0 = pto.alloc_tile addr = %a1536 : !pto.tile_buf + %re_src1 = pto.alloc_tile addr = %a1792 : !pto.tile_buf + %re_dst = pto.alloc_tile addr = %a2048 : !pto.tile_buf + pto.trowexpandadd ins(%re_src0, %re_src1 : !pto.tile_buf, !pto.tile_buf) + outs(%re_dst : !pto.tile_buf) + return + } +} + +// A5-LABEL: func.func @a5_skip_no_tmp_overloads +// A5: pto.tci ins(%{{.*}} : i32) outs( +// A5: pto.tcvt ins(%{{.*}} {{.*}} : !pto.tile_buf) outs(%{{.*}} : !pto.tile_buf) +// A5: pto.tquant ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.trowexpandadd ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) diff --git a/test/lit/pto/implicit_tmp_arg_reductions.pto b/test/lit/pto/implicit_tmp_arg_reductions.pto index 4f51f1da7c..b144227d33 100644 --- a/test/lit/pto/implicit_tmp_arg_reductions.pto +++ b/test/lit/pto/implicit_tmp_arg_reductions.pto @@ -6,8 +6,8 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR -// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 module { func.func @implicit_arg_reduction_tmps() { @@ -24,8 +24,13 @@ module { } } -// IR-LABEL: func.func @implicit_arg_reduction_tmps -// IR: pto.tcolargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) -// IR: pto.tcolargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) -// IR: pto.trowargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) -// IR: pto.trowargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3-LABEL: func.func @implicit_arg_reduction_tmps +// A3: pto.tcolargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3: pto.tcolargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3: pto.trowargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3: pto.trowargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5-LABEL: func.func @implicit_arg_reduction_tmps +// A5: pto.tcolargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.tcolargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.trowargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.trowargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) diff --git a/test/lit/pto/implicit_tmp_row_reductions.pto b/test/lit/pto/implicit_tmp_row_reductions.pto index cfd3ade699..15c62e65b4 100644 --- a/test/lit/pto/implicit_tmp_row_reductions.pto +++ b/test/lit/pto/implicit_tmp_row_reductions.pto @@ -6,8 +6,8 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR -// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 module { func.func @implicit_row_reduction_tmps() { @@ -24,8 +24,13 @@ module { } } -// IR-LABEL: func.func @implicit_row_reduction_tmps -// IR: pto.trowmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) -// IR: pto.trowmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) -// IR: pto.trowsum ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) -// IR: pto.trowprod ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3-LABEL: func.func @implicit_row_reduction_tmps +// A3: pto.trowmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3: pto.trowmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3: pto.trowsum ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3: pto.trowprod ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5-LABEL: func.func @implicit_row_reduction_tmps +// A5: pto.trowmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.trowmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.trowsum ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.trowprod ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) diff --git a/test/lit/pto/implicit_tmp_xor_materialization.pto b/test/lit/pto/implicit_tmp_xor_materialization.pto index 172d428c80..d5e27a6dc4 100644 --- a/test/lit/pto/implicit_tmp_xor_materialization.pto +++ b/test/lit/pto/implicit_tmp_xor_materialization.pto @@ -6,8 +6,8 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR -// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 module { func.func @implicit_xor_tmps() { @@ -22,6 +22,9 @@ module { } } -// IR-LABEL: func.func @implicit_xor_tmps -// IR: pto.txor ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) -// IR: pto.txors ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, i16, !pto.tile_buf) +// A3-LABEL: func.func @implicit_xor_tmps +// A3: pto.txor ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) +// A3: pto.txors ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, i16, !pto.tile_buf) +// A5-LABEL: func.func @implicit_xor_tmps +// A5: pto.txor ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) +// A5: pto.txors ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, i16, !pto.tile_buf) diff --git a/test/lit/pto/plan_memory_inplace_forbid_alias.pto b/test/lit/pto/plan_memory_inplace_forbid_alias.pto index 92062d448e..8bfed61303 100644 --- a/test/lit/pto/plan_memory_inplace_forbid_alias.pto +++ b/test/lit/pto/plan_memory_inplace_forbid_alias.pto @@ -26,14 +26,14 @@ module attributes {"pto.target_arch" = "a3"} { %mask = pto.alloc_tile : !pto.tile_buf %src0 = pto.alloc_tile : !pto.tile_buf %src1 = pto.alloc_tile : !pto.tile_buf - %tmp = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, - !pto.tile_buf) + !pto.tile_buf) outs(%dst : !pto.tile_buf) return } diff --git a/test/lit/pto/select_tile_native.pto b/test/lit/pto/select_tile_native.pto index a0ba395ba0..9e428bce07 100644 --- a/test/lit/pto/select_tile_native.pto +++ b/test/lit/pto/select_tile_native.pto @@ -10,8 +10,8 @@ // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC module { - func.func private @tsel_arg(%mask: !pto.tile_buf, %a: !pto.tile_buf, %b: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { - pto.tsel ins(%mask, %a, %b, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + func.func private @tsel_arg(%mask: !pto.tile_buf, %a: !pto.tile_buf, %b: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tsel ins(%mask, %a, %b, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) return } func.func private @tsels_arg(%mask: !pto.tile_buf, %src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf, %scalar: i16) { diff --git a/test/lit/pto/tsel_bf16.pto b/test/lit/pto/tsel_bf16.pto index f0dea4d774..cbc41cc577 100644 --- a/test/lit/pto/tsel_bf16.pto +++ b/test/lit/pto/tsel_bf16.pto @@ -14,10 +14,10 @@ module { %mask = pto.alloc_tile : !pto.tile_buf %src0 = pto.alloc_tile : !pto.tile_buf %src1 = pto.alloc_tile : !pto.tile_buf - %tmp = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) return } diff --git a/test/lit/pto/tsel_tmp_contract_a3_invalid.pto b/test/lit/pto/tsel_tmp_contract_a3_invalid.pto new file mode 100644 index 0000000000..29b79e0ac4 --- /dev/null +++ b/test/lit/pto/tsel_tmp_contract_a3_invalid.pto @@ -0,0 +1,24 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s + +module attributes {"pto.device-spec" = "Ascend910B3"} { + func.func @a3_tsel_tmp_too_small() attributes {pto.kernel_kind = #pto.kernel_kind} { + %mask = pto.alloc_tile : !pto.tile_buf + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK: error: 'pto.tsel' op expects tmp capacity to be at least 64 bytes, but got 32 bytes From 24da54612781003266cb2f30cda54d01b90ade7e Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 7 Aug 2026 09:32:44 +0800 Subject: [PATCH 055/122] Relax tsel tmp verifier --- lib/PTO/IR/PTO.cpp | 5 ++--- test/lit/pto/tsel_tmp_contract_a3_4byte.pto | 25 +++++++++++++++++++++ 2 files changed, 27 insertions(+), 3 deletions(-) create mode 100644 test/lit/pto/tsel_tmp_contract_a3_4byte.pto diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 7c1f269e72..fa24568abc 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -13072,9 +13072,8 @@ mlir::LogicalResult mlir::pto::TSelOp::verify() { "expects A2/A3 tsel src0, src1, and dst element type to be i16/i32/f16/bf16/f32"); if (getTmp()) { Type tmpTy = getTmp().getType(); - auto tmpElem = dyn_cast(getElemTy(tmpTy)); - if (!tmpElem || tmpElem.getWidth() != 32) - return emitOpError("expects A2/A3 tsel tmp element type to be i32"); + if (getElemByteSize(getElemTy(tmpTy)) != 4) + return emitOpError("expects A2/A3 tsel tmp element type to be 4 bytes wide"); unsigned elemBits = getPTOStorageElemBitWidth(elem); if (elemBits != 16 && elemBits != 32) return emitOpError("expects A2/A3 tsel data element type to be 16 or 32 bits"); diff --git a/test/lit/pto/tsel_tmp_contract_a3_4byte.pto b/test/lit/pto/tsel_tmp_contract_a3_4byte.pto new file mode 100644 index 0000000000..09c894e950 --- /dev/null +++ b/test/lit/pto/tsel_tmp_contract_a3_4byte.pto @@ -0,0 +1,25 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s + +module attributes {"pto.device-spec" = "Ascend910B3"} { + func.func @a3_tsel_f32_tmp() attributes {pto.kernel_kind = #pto.kernel_kind} { + %mask = pto.alloc_tile : !pto.tile_buf + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK-LABEL: func.func @a3_tsel_f32_tmp +// CHECK: pto.tsel From 3df5f517ca8b01fa5f341b0071208bbba21746e4 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 7 Aug 2026 09:40:38 +0800 Subject: [PATCH 056/122] Remove redundant tsel tmp lit --- test/lit/pto/tsel_tmp_contract_a3_4byte.pto | 25 --------------------- 1 file changed, 25 deletions(-) delete mode 100644 test/lit/pto/tsel_tmp_contract_a3_4byte.pto diff --git a/test/lit/pto/tsel_tmp_contract_a3_4byte.pto b/test/lit/pto/tsel_tmp_contract_a3_4byte.pto deleted file mode 100644 index 09c894e950..0000000000 --- a/test/lit/pto/tsel_tmp_contract_a3_4byte.pto +++ /dev/null @@ -1,25 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -// RUN: ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s - -module attributes {"pto.device-spec" = "Ascend910B3"} { - func.func @a3_tsel_f32_tmp() attributes {pto.kernel_kind = #pto.kernel_kind} { - %mask = pto.alloc_tile : !pto.tile_buf - %src0 = pto.alloc_tile : !pto.tile_buf - %src1 = pto.alloc_tile : !pto.tile_buf - %tmp = pto.alloc_tile : !pto.tile_buf - %dst = pto.alloc_tile : !pto.tile_buf - pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) - outs(%dst : !pto.tile_buf) - return - } -} - -// CHECK-LABEL: func.func @a3_tsel_f32_tmp -// CHECK: pto.tsel From f9980f70c40c74b4a5bc69f74798e005f86ad411 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 7 Aug 2026 11:44:09 +0800 Subject: [PATCH 057/122] Tighten tsel tmp capacity to backend usage --- lib/PTO/IR/PTO.cpp | 2 +- test/lit/pto/tsel_tmp_contract_a3_invalid.pto | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index fa24568abc..055d53f740 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -13077,7 +13077,7 @@ mlir::LogicalResult mlir::pto::TSelOp::verify() { unsigned elemBits = getPTOStorageElemBitWidth(elem); if (elemBits != 16 && elemBits != 32) return emitOpError("expects A2/A3 tsel data element type to be 16 or 32 bits"); - uint64_t minBytes = elemBits == 16 ? 64 : 32; + uint64_t minBytes = elemBits == 16 ? 16 : 8; if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, minBytes))) return failure(); } diff --git a/test/lit/pto/tsel_tmp_contract_a3_invalid.pto b/test/lit/pto/tsel_tmp_contract_a3_invalid.pto index 29b79e0ac4..937c2d7e35 100644 --- a/test/lit/pto/tsel_tmp_contract_a3_invalid.pto +++ b/test/lit/pto/tsel_tmp_contract_a3_invalid.pto @@ -9,16 +9,15 @@ // RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s module attributes {"pto.device-spec" = "Ascend910B3"} { - func.func @a3_tsel_tmp_too_small() attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @a3_tsel_tmp_too_small(%tmp: !pto.tile_buf) attributes {pto.kernel_kind = #pto.kernel_kind} { %mask = pto.alloc_tile : !pto.tile_buf %src0 = pto.alloc_tile : !pto.tile_buf %src1 = pto.alloc_tile : !pto.tile_buf - %tmp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) return } } -// CHECK: error: 'pto.tsel' op expects tmp capacity to be at least 64 bytes, but got 32 bytes +// CHECK: error: 'pto.tsel' op expects tmp capacity to be at least 16 bytes, but got 8 bytes From 51f858fb269e411a169eccf9d8e82b082601c33c Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 7 Aug 2026 16:08:23 +0800 Subject: [PATCH 058/122] Refresh implicit-tmp design for capacity and A5 level3 Align the doc with actual materialize/verifier behavior: reduction tmp capacity is checked from declared shape and the 32B floor is always satisfied because alloc_tile enforces 32B row alignment; A5 level3 synthesizes a no-effect 32B ABI placeholder instead of erroring. --- ...oas-implicit-tmp-materialization-design.md | 47 +++++++++++++------ 1 file changed, 32 insertions(+), 15 deletions(-) diff --git a/docs/designs/ptoas-implicit-tmp-materialization-design.md b/docs/designs/ptoas-implicit-tmp-materialization-design.md index aaad9a2276..8904bd9883 100644 --- a/docs/designs/ptoas-implicit-tmp-materialization-design.md +++ b/docs/designs/ptoas-implicit-tmp-materialization-design.md @@ -114,22 +114,31 @@ target_op(no tmp) ### level3 -level3 下用户显式管理 local 地址，memplan 通常跳过。因此不应自动创建无地址 tmp。 +level3 下用户显式管理 local 地址，memplan 通常跳过。pass 通过构造参数 `requireExplicitTmp` 判定 level3（`createPTOMaterializeImplicitTmpPass(effectiveLevel == Level3)`）。 -通用规则： +对 **A2/A3 实际使用 tmp 的 op**，level3 不自动创建无地址 tmp，缺省 tmp 直接报错： ```text -level3 + target_op(no tmp) => pass/verifier 报错 +level3 + A2/A3 target_op(no tmp) => pass 报错 ``` -用户在 level3 使用已纳入改造的目标 op 时，必须显式提供合法 tmp，并保证 tmp 自身带合法 local addr，或满足现有 level3 显式地址规则。 - -诊断信息示例： +实际诊断字符串（`PTOMaterializeImplicitTmp.cpp`）： ```text - requires explicit tmp when compiling at level3 because PlanMemory is skipped + requires explicit tmp when PlanMemory is skipped ``` +个别 op 有更具体的变体，例如 binary tcolsum 为 `requires explicit tmp for binary tcolsum when PlanMemory is skipped`，非 32 对齐 tsort32 为 `requires explicit tmp for non-32-aligned tsort32 when PlanMemory is skipped`。 + +A2/A3 用户在 level3 使用这些 op 时，必须显式提供合法 tmp，并保证 tmp 自身带合法 local addr，或满足现有 level3 显式地址规则。 + +**A5 例外**：对 A5 只接受但不使用 tmp 的 op，pass 通过 `!isA5` 跳过上述 level3 报错： + +- 若后端 C++ 签名仍要求 tmp（row/arg reduction、TXOR/TXORS、TSEL/TSELS、TPRELU、TREM/TREMS 等），pass 即使在 level3 也自动生成固定 32 字节的 ABI placeholder（`makeA5PlaceholderTmpType`，形状 `{1, 32/sizeof(elem)}`）。该 placeholder 不建模 Read/Write MemoryEffects、不参与 memplan、无需回写地址，因此不违反 level3“不自动分配 tmp 地址”的非目标。 +- 若后端存在 no-tmp overload（TCI、TROWEXPAND*、TQUANT 等），A5 直接保持 no-tmp 形态，不补 placeholder。 + +因此当前实现下不存在 “A5 + level3 + 缺省 tmp” 报错的路径：要么 pass 自动补 placeholder，要么保持合法的 no-tmp 形态，verifier 的 no-tmp overload 也无条件接受。 + ## EmitC Lowering 目标 op 的 EmitC lowering 应保持简单： @@ -594,10 +603,11 @@ TSORT32, TQUANT ### 通用规则 - level1/level2：只有该 op 在当前 arch / 模式下实际需要 tmp，且 IR 允许省略 tmp 时，才自动 materialize `pto.alloc_tile(no addr)`。 -- level3：若该 op 在当前 arch / 模式下需要 tmp 且用户省略 tmp，则报错；若当前 arch / 模式不使用 tmp，则不强制补 tmp。 -- A5 仅接受但不使用 tmp 的 op：若后端存在 no-tmp overload，则不自动补 tmp；若 C++ 签名仍要求 tmp，则自动生成 ABI placeholder。placeholder 不建模 tmp 的 Read/Write，用户显式 tmp 也不按 A2/A3 容量规则校验。 +- level3：若该 op 在 A2/A3 当前模式下需要 tmp 且用户省略 tmp，则报错（`requires explicit tmp when PlanMemory is skipped`）。A5 见下一条，即使 level3 也不因缺省 tmp 报错。 +- A5 仅接受但不使用 tmp 的 op：若后端存在 no-tmp overload，则不自动补 tmp；若 C++ 签名仍要求 tmp，则自动生成固定 32 字节 ABI placeholder（level1/2/3 一致，通过 `!isA5` 绕过 level3 显式 tmp 检查）。placeholder 不建模 tmp 的 Read/Write（`getEffects` 以 `!tmp.empty() && arch != A5` 守卫）、不参与 memplan、无需地址，用户显式 tmp 也不按 A2/A3 容量规则校验。 - tmp 是 scratch 的 op：MemoryEffects 需要建模为 `Read(tmp) + Write(tmp)`，或在 semantic no-alias side table 中显式加入 `forbidAlias(tmp, dst/output)`。 - 原 mandatory tmp op 已统一改为 optional，并保持显式 tmp 文本格式兼容。 +- 容量校验统一走 `verifyTmpCapacityAtLeast`，按 tmp 的**声明 shape** × `sizeof(dtype)` 计算（`getStaticByteSize`，非 valid 区域），A5 分支不执行该校验。对 row/arg reduction 的 32 字节下限：因 `pto.alloc_tile` 已对 row-major none_box tile 强制行 `cols * sizeof(dtype)` 32 字节对齐，而 reduction `src` 必须是这类 tile，故任何合法 src 单行即 ≥32 字节，同形状 tmp 恒满足下限；materialize 无需为 sub-block src 额外兜底容量（<32 字节的合法 reduction src 无法构造）。 ### Arg reduction 类 @@ -625,9 +635,15 @@ TROWARGMAX / TROWARGMIN： - tmp 行数与 `src` 相同；每行 stride 按 PTO-ISA 文档公式计算。 - 当前 PTOAS ODS 只有单输出索引模式；该模式仍需满足后端显式 tmp 参数签名，因此 level1/2 生成保守同形状 tmp。未来接入值+索引模式后，再按输出模式和归约阶段收紧容量。 +容量校验： + +- A2/A3 verifier 在按 `tmpGapEles` / 布局（DN 1 列、ND 2 列、min stride 等）校验后，统一以 `verifyTmpCapacityAtLeast(op, tmp, 32)` 兜底；容量按声明 shape 计算，同 row reduction 一样被 `pto.alloc_tile` 的 32 字节行对齐恒满足。 +- A5 arg-reduction verifier（`verifyTColArgReductionOpA5` / `verifyTRowArgReductionOpA5`）不含任何容量校验。 + MemoryEffects / alias： - A2/A3 实际使用 tmp 时建模 `Read(src) + Read(tmp) + Write(tmp) + Write(dstIdx/dstVal)`。 +- A5 placeholder 不建模 tmp 的 Read/Write（`getEffects` 以 `!tmp.empty() && arch != A5` 守卫）。 - tmp 不应与同 op 的输出 alias；如果 MemoryEffects 无法覆盖，应加入 `forbidAlias(tmp, dstIdx)` 和必要的 `forbidAlias(tmp, dstVal)`。 ### Row reduction 类 @@ -646,19 +662,20 @@ TROWPROD, TROWSUM, TROWMAX, TROWMIN tmp 规格： - tmp dtype 与 `src` / `dst` 一致。 -- 整数路径最小需要 1 个 vector block：`int32` 为 8 列，`int16` 为 16 列。 -- 浮点路径用于二叉树归约；安全默认形状可设为与 `src` 相同。 -- `TROWPROD` 的安全默认也可设为与 `src` 相同；最小需求为 1 行和 1 个 vector block。 +- ISA 最小需求为 1 行 1 个 vector block（32 字节）：`int32` 为 8 列，`int16` 为 16 列；浮点二叉树归约同样以 1 个 block 为下限。 +- A2/A3 materialize 直接生成与 `src` 同形状的 tmp（`makeSameShapeTmpType`），不做逐 dtype 的特化裁剪；`TROWPROD` 亦同。 +- 容量校验：A2/A3 verifier 走 `verifyTmpCapacityAtLeast(op, tmp, 32)`，按声明 shape × `sizeof(dtype)` 计算。**该 32 字节下限对合法 IR 恒被满足**——reduction `src` 必须是 row-major none_box tile，`pto.alloc_tile` 已对其强制行 32 字节对齐，故同形状 tmp 单行即 ≥32 字节，无需为 sub-block src 特殊兜底。 Pass 行为： -- A2/A3 level1/2：若 IR 已支持 optional tmp 且缺省 tmp，则自动生成 vec row-major none-box tmp。 -- A5：生成后端签名需要的 ABI placeholder，但不为其添加 tmp MemoryEffects。 -- level3：A2/A3 需要 tmp 时缺省 tmp 报错；A5 不强制 tmp。 +- A2/A3 level1/2：若 IR 已支持 optional tmp 且缺省 tmp，则自动生成与 `src` 同形状的 vec row-major none-box tmp。 +- A5：生成后端签名需要的固定 32 字节 ABI placeholder，但不为其添加 tmp MemoryEffects，也不跑 32 字节容量校验。 +- level3：A2/A3 需要 tmp 时缺省 tmp 报错（`requires explicit tmp when PlanMemory is skipped`）；A5 通过 `!isA5` 跳过报错，仍自动生成 ABI placeholder。 MemoryEffects / alias： - A2/A3 建模 `Read(src) + Read(tmp) + Write(tmp) + Write(dst)`。 +- A5 placeholder 不建模 tmp 的 Read/Write（`getEffects` 以 `!tmp.empty() && arch != A5` 守卫）。 - tmp 与 `dst` 禁止 alias。 ### Column sum 类 From db9d6e47d9b8ea77ec599675fa63bf87efa9c1c4 Mon Sep 17 00:00:00 2001 From: and0d0 Date: Fri, 7 Aug 2026 16:46:41 +0800 Subject: [PATCH 059/122] fix(ptodsl): guard fp4 explicit l1 to l0 loads --- ptodsl/ptodsl/_ops.py | 36 +++++++++++++++++---- ptodsl/tests/test_vector_cube_ops.py | 19 +++++++++++ test/lit/pto/mte_l1_l0_fp4_s4_expand_a5.pto | 4 +++ 3 files changed, 53 insertions(+), 6 deletions(-) diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 9479563605..3ae177fd83 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -5276,6 +5276,22 @@ def mem_bar(barrier_type): _pto.MemBarOp(kind=_membar_attr(barrier_name)) +def _is_fp4_packed_pointer_value(value) -> bool: + type_text = str(getattr(value, "type", "")) + return type_text.startswith("!pto.ptr") + destination = object() + controls = { + "m_start": 3, + "k_start": 5, + "m_step": 16, + "k_step": 2, + "src_stride": 8, + "dst_stride": 2, + } + + with patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "_coerce_i64", side_effect=lambda value, *, context: f"{context}:{value}"): + with self.assertRaisesRegex(TypeError, "explicit-control FP4 loads are not supported yet"): + _ops.mte_l1_l0a(source, destination, **controls, transpose=True) + with self.assertRaisesRegex(TypeError, "using FP4 in source may silently select an incorrect intrinsic"): + _ops.mte_l1_l0b(source, destination, **controls, transpose=True) + def test_mte_l1_l0_legacy_forms_preserve_keyword_compatibility(self): source = object() destination = object() diff --git a/test/lit/pto/mte_l1_l0_fp4_s4_expand_a5.pto b/test/lit/pto/mte_l1_l0_fp4_s4_expand_a5.pto index 11dd025eb1..fe3a2bd4c0 100644 --- a/test/lit/pto/mte_l1_l0_fp4_s4_expand_a5.pto +++ b/test/lit/pto/mte_l1_l0_fp4_s4_expand_a5.pto @@ -10,6 +10,7 @@ // bridge ops, with K coordinates converted to s4 units instead of byte-path // defaults. // RUN: ptoas --pto-arch=a5 --pto-level=level3 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s +// RUN: ( mkdir -p %T && ptoas --pto-arch=a5 --pto-level=level3 --pto-backend=vpto %s -o %t --mlir-print-ir-after=convert-func-to-llvm 2>&1 || true ) | FileCheck %s --check-prefix=LLVM module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { func.func @fp4_mte_l1_l0_expand( @@ -30,3 +31,6 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind Date: Fri, 7 Aug 2026 16:47:28 +0800 Subject: [PATCH 060/122] fix(vpto): validate cube bridge load control widths --- lib/PTO/IR/VPTO.cpp | 38 +++++++++++-------- .../vpto/load_cbuf_to_l0_verifier_invalid.pto | 30 +++++++++++++++ 2 files changed, 52 insertions(+), 16 deletions(-) diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index 591c3cc8d5..623c3fb329 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -7806,26 +7806,32 @@ static LogicalResult verifyMxDestinationAlignment(Operation *op, template static LogicalResult verifyExplicitCubeBridgeLoadControls(OpTy op) { - auto checkNonNegativeConst = [&](Value value, StringRef name) -> LogicalResult { - APInt intValue; - if (matchPattern(value, m_ConstantInt(&intValue)) && intValue.isNegative()) - return op.emitOpError() << name << " must be non-negative"; - return success(); - }; - auto checkPositiveConst = [&](Value value, StringRef name) -> LogicalResult { + auto checkConstRange = [&](Value value, StringRef name, int64_t min, + int64_t max) -> LogicalResult { APInt intValue; - if (matchPattern(value, m_ConstantInt(&intValue)) && - (intValue.isNegative() || intValue.isZero())) - return op.emitOpError() << name << " must be greater than zero"; + if (!matchPattern(value, m_ConstantInt(&intValue))) + return success(); + int64_t signedValue = intValue.getSExtValue(); + if (signedValue < min) + return op.emitOpError() + << name + << (min == 0 ? " must be non-negative" + : " must be greater than zero"); + if (signedValue > max) + return op.emitOpError() + << name << " must be <= " << max + << " to fit the hardware control field"; return success(); }; - if (failed(checkNonNegativeConst(op.getMStart(), "m_start")) || - failed(checkNonNegativeConst(op.getKStart(), "k_start")) || - failed(checkPositiveConst(op.getMStep(), "m_step")) || - failed(checkPositiveConst(op.getKStep(), "k_step")) || - failed(checkPositiveConst(op.getSrcStride(), "src_stride")) || - failed(checkPositiveConst(op.getDstStride(), "dst_stride"))) + constexpr int64_t kU16Max = 65535; + constexpr int64_t kU8Max = 255; + if (failed(checkConstRange(op.getMStart(), "m_start", 0, kU16Max)) || + failed(checkConstRange(op.getKStart(), "k_start", 0, kU16Max)) || + failed(checkConstRange(op.getMStep(), "m_step", 1, kU8Max)) || + failed(checkConstRange(op.getKStep(), "k_step", 1, kU8Max)) || + failed(checkConstRange(op.getSrcStride(), "src_stride", 1, kU16Max)) || + failed(checkConstRange(op.getDstStride(), "dst_stride", 1, kU16Max))) return failure(); return success(); } diff --git a/test/lit/vpto/load_cbuf_to_l0_verifier_invalid.pto b/test/lit/vpto/load_cbuf_to_l0_verifier_invalid.pto index 2145833866..c959568133 100644 --- a/test/lit/vpto/load_cbuf_to_l0_verifier_invalid.pto +++ b/test/lit/vpto/load_cbuf_to_l0_verifier_invalid.pto @@ -11,6 +11,8 @@ // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_destination.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-DESTINATION // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/cb_negative_m_start.pto -o - 2>&1 | FileCheck %s --check-prefix=CB-NEGATIVE-M-START // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/cb_zero_m_step.pto -o - 2>&1 | FileCheck %s --check-prefix=CB-ZERO-M-STEP +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_m_step_overflow.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-M-STEP-OVERFLOW +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/ca_m_start_overflow.pto -o - 2>&1 | FileCheck %s --check-prefix=CA-M-START-OVERFLOW //--- ca_source.pto module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { @@ -61,7 +63,35 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @invalid_control_values() attributes {pto.kernel} { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c256 = arith.constant 256 : i64 + %src = pto.castptr %c0 : i64 -> !pto.ptr + %dst = pto.castptr %c0 : i64 -> !pto.ptr + pto.load_cbuf_to_ca %src, %dst, %c0, %c0, %c256, %c1, %c1, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + return + } +} + +//--- ca_m_start_overflow.pto +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @invalid_control_values() attributes {pto.kernel} { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c65536 = arith.constant 65536 : i64 + %src = pto.castptr %c0 : i64 -> !pto.ptr + %dst = pto.castptr %c0 : i64 -> !pto.ptr + pto.load_cbuf_to_ca %src, %dst, %c65536, %c0, %c1, %c1, %c1, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + return + } +} + // CA-SOURCE: 'pto.load_cbuf_to_ca' op requires MAT source // CA-DESTINATION: 'pto.load_cbuf_to_ca' op requires LEFT destination // CB-NEGATIVE-M-START: 'pto.load_cbuf_to_cb' op m_start must be non-negative // CB-ZERO-M-STEP: 'pto.load_cbuf_to_cb' op m_step must be greater than zero +// CA-M-STEP-OVERFLOW: 'pto.load_cbuf_to_ca' op m_step must be <= 255 to fit the hardware control field +// CA-M-START-OVERFLOW: 'pto.load_cbuf_to_ca' op m_start must be <= 65535 to fit the hardware control field From 3c3036f462c5c74c25d1a4277f2016a8de02ce43 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Fri, 7 Aug 2026 17:01:01 +0800 Subject: [PATCH 061/122] fix(ptodsl): preserve sequential conditional values (#1183) --- ptodsl/ptodsl/_ast_rewrite.py | 24 ++++++++++++--- ptodsl/tests/test_section.py | 57 +++++++++++++++++++++++++++++++++++ 2 files changed, 76 insertions(+), 5 deletions(-) diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index 787ee0d83d..ee416cbf60 100644 --- a/ptodsl/ptodsl/_ast_rewrite.py +++ b/ptodsl/ptodsl/_ast_rewrite.py @@ -140,11 +140,25 @@ def _visit_section_body(self, stmts): old_env = self._env old_names = self._local_names old_outer_bindings = self._section_outer_bindings + entry_binding_count = len(self.section_entry_bindings) self._env = {} self._local_names = _name_info(stmts).stores self._section_outer_bindings = set(self._known_bindings) try: - return [self.visit(stmt) for stmt in stmts] + body = [self.visit(stmt) for stmt in stmts] + # Materialize outer values under their section-local aliases before + # any runtime control flow. Subsequent branch merges can then read + # the alias at the current program point instead of always falling + # back to the section entry value. + entry_bindings = list(self.section_entry_bindings.items())[entry_binding_count:] + initializers = [ + ast.Assign( + targets=[_name(alias, ast.Store())], + value=_name(outer_name), + ) + for alias, outer_name in entry_bindings + ] + return initializers + body finally: self._env = old_env self._local_names = old_names @@ -970,13 +984,13 @@ def _fresh(self, prefix: str) -> str: self._counter += 1 return value - def _section_entry_value(self, name): + def _current_value(self, name): if name in self._section_uninitialized_aliases: raise PTODSLAstRewriteError( "ast_rewrite=True runtime if reads a section-local value before it is initialized; " f"initialize {name!r} before the conditional" ) - return _name(self._section_entry_bindings.get(name, name)) + return _name(name) def rewrite_block(self, stmts, *, live_after, live_after_slots=None, allow_loop_control=False, static_iters=None): rewritten_reversed = [] @@ -1232,7 +1246,7 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con result.extend( ast.Assign( targets=[_name(old_name, ast.Store())], - value=self._section_entry_value(name), + value=self._current_value(name), ) for name, old_name in old_value_names.items() ) @@ -1372,7 +1386,7 @@ def _rewrite_for(self, stmt, *, live_after, live_after_slots=None, allow_loop_co keywords=[ ast.keyword( arg=name, - value=_name(self._section_entry_bindings.get(name, name)), + value=self._current_value(name), ) for name in loop_carried ] + [ diff --git a/ptodsl/tests/test_section.py b/ptodsl/tests/test_section.py index 5e7b132010..517eb8885f 100644 --- a/ptodsl/tests/test_section.py +++ b/ptodsl/tests/test_section.py @@ -9,6 +9,8 @@ """Focused tracing coverage for explicit physical section hints.""" +import re + from ptodsl import pto from ptodsl._ast_rewrite import PTODSLAstRewriteError from ptodsl._context import make_context @@ -199,6 +201,34 @@ def lexical_section_sibling_single_sided_conditional_rebinding_probe(): pto.wait_flag("MTE2", "S", event_id=n_tile_2) +@pto.jit(target="a5", mode="explicit") +def lexical_section_sequential_single_sided_conditional_probe(): + m_tile = pto.const(0, dtype=pto.i64) + n_tile = pto.const(0, dtype=pto.i64) + with pto.section("cube"): + if pto.get_block_idx() < 16: + m_tile = pto.get_block_idx() & 3 + n_tile = pto.get_block_idx() // 4 + if 16 <= pto.get_block_idx(): + m_tile = (pto.get_block_idx() & 3) + 4 + n_tile = (pto.get_block_idx() // 4) - 4 + if 15 < pto.get_block_idx(): + n_tile = 3 - n_tile + pto.wait_flag("S", "MTE2", event_id=m_tile + n_tile) + + +@pto.jit(target="a5", mode="explicit") +def lexical_section_single_sided_read_before_rebinding_probe(): + value = pto.const(0, dtype=pto.i64) + with pto.section("cube"): + if pto.get_block_idx() < 16: + previous_value = value + value = pto.get_block_idx() + else: + previous_value = value + pto.wait_flag("S", "MTE2", event_id=previous_value + value) + + @pto.jit(target="a5", mode="explicit") def lexical_section_uninitialized_conditional_probe(): one = pto.const(1, dtype=pto.i32) @@ -342,6 +372,33 @@ def main() -> None: module = Module.parse(sibling_single_sided_text, context) module.operation.verify() + sequential_single_sided_text = lexical_section_sequential_single_sided_conditional_probe.compile().mlir_text() + if_results = re.findall(r"^\s*(%\d+)(?::\d+)? = scf\.if", sequential_single_sided_text, re.MULTILINE) + assert len(if_results) == 3 + second_if_text = sequential_single_sided_text.split(f"{if_results[1]}:2 = scf.if", 1)[1] + second_if_text = second_if_text.split(f"{if_results[2]} = scf.if", 1)[0] + assert re.search( + rf"else \{{\s+scf\.yield {re.escape(if_results[0])}#0, {re.escape(if_results[0])}#1 : i64, i64", + second_if_text, + ) + third_if_text = sequential_single_sided_text.split(f"{if_results[2]} = scf.if", 1)[1] + assert re.search( + rf"else \{{\s+scf\.yield {re.escape(if_results[1])}#1 : i64", + third_if_text, + ) + with make_context() as context: + module = Module.parse(sequential_single_sided_text, context) + module.operation.verify() + + read_before_rebinding_text = lexical_section_single_sided_read_before_rebinding_probe.compile().mlir_text() + assert re.search( + r"scf\.yield %c0_i64, %[\d]+ : i64, i64", + read_before_rebinding_text, + ) + with make_context() as context: + module = Module.parse(read_before_rebinding_text, context) + module.operation.verify() + nested_conditional_text = lexical_section_nested_conditional_rebinding_probe.compile().mlir_text() assert nested_conditional_text.count("pto.section.cube {") == 1 assert nested_conditional_text.count("scf.if") == 2 From cee6206f91ae6170391abae59c87c483ac857eb4 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Mon, 3 Aug 2026 10:04:59 +0800 Subject: [PATCH 062/122] feat(tilelib): materialize PTODSL in process --- docker/test_wheel_imports.sh | 9 - .../ptoas-compiler-dso-wheel-linking.md | 3 +- docs/designs/ptoas-python-launcher-layout.md | 14 +- ...todsl-tilelib-template-selection-design.md | 26 +- include/PTO/Transforms/Passes.h | 6 +- include/PTO/Transforms/Passes.td | 51 +-- include/PTO/Transforms/TileLibService.h | 54 +++ lib/PTO/Transforms/ExpandTileOp.cpp | 311 +++++++----------- .../Transforms/InsertTemplateAttributes.cpp | 162 ++------- ptodsl/README.md | 11 +- .../tilelib-debugging-playbook.md | 18 +- .../tilelib-template-authoring.md | 2 +- ptodsl/ptoas/_cli.py | 30 +- ptodsl/ptodsl/tilelib/_compiler_runtime.py | 72 ++++ ptodsl/ptodsl/tilelib/_render_runtime.py | 4 +- .../{serving/daemon.py => _selection.py} | 218 +----------- ptodsl/ptodsl/tilelib/decorator.py | 13 + ptodsl/ptodsl/tilelib/serving/__init__.py | 33 -- ptodsl/ptodsl/tilelib/serving/client.py | 77 ----- ptodsl/ptodsl/tilelib/serving/helper.py | 73 ---- ptodsl/ptodsl/tilelib/serving/wire.py | 57 ---- ptodsl/tests/test_ptoas_cli.py | 24 +- ptodsl/tests/test_tilelib_daemon.py | 266 --------------- ptodsl/tests/test_tilelib_render.py | 18 + test/lit/vpto/expand_tile_op_ptodsl_tsub.pto | 6 +- ...xpand_tile_op_ptodsl_view_stride_cache.pto | 2 +- .../st/smoke/testcase/run_ptoas_to_file.cmake | 11 - .../src/st/testcase/run_ptoas_to_file.cmake | 11 - .../script/run_a5_st_all_parallel.py | 20 +- .../script/run_ptodsl_st_parallel.py | 18 +- tools/ptoas/CMakeLists.txt | 8 - tools/ptoas/NativeModule.cpp | 101 +++++- tools/ptoas/TilelangDaemon.cpp | 154 --------- tools/ptoas/TilelangDaemon.h | 44 --- tools/ptoas/driver.cpp | 65 +++- tools/ptoas/ptoas.cpp | 211 +----------- tools/ptoas/ptoas.h | 12 +- 37 files changed, 547 insertions(+), 1668 deletions(-) create mode 100644 include/PTO/Transforms/TileLibService.h create mode 100644 ptodsl/ptodsl/tilelib/_compiler_runtime.py rename ptodsl/ptodsl/tilelib/{serving/daemon.py => _selection.py} (58%) delete mode 100644 ptodsl/ptodsl/tilelib/serving/__init__.py delete mode 100644 ptodsl/ptodsl/tilelib/serving/client.py delete mode 100644 ptodsl/ptodsl/tilelib/serving/helper.py delete mode 100644 ptodsl/ptodsl/tilelib/serving/wire.py delete mode 100644 ptodsl/tests/test_tilelib_daemon.py delete mode 100644 tools/ptoas/TilelangDaemon.cpp delete mode 100644 tools/ptoas/TilelangDaemon.h diff --git a/docker/test_wheel_imports.sh b/docker/test_wheel_imports.sh index 85a85aa143..4c1a0a318f 100755 --- a/docker/test_wheel_imports.sh +++ b/docker/test_wheel_imports.sh @@ -222,13 +222,4 @@ grep -q "candidates = " "${CLEAN_ENV_PTO_IR}" || { echo "Error: clean-environment ptoas smoke output is missing TileLib candidate metadata" >&2 exit 1 } -if ! grep -q "TileLib daemon started successfully" "${CLEAN_ENV_LOG}"; then - echo "Error: TileLib daemon did not report a successful start" >&2 - exit 1 -fi -if ! grep -q "TileLib daemon stopped" "${CLEAN_ENV_LOG}"; then - echo "Error: TileLib daemon did not report a clean stop" >&2 - exit 1 -fi - echo "All wheel import tests passed!" diff --git a/docs/designs/ptoas-compiler-dso-wheel-linking.md b/docs/designs/ptoas-compiler-dso-wheel-linking.md index bdd3bed827..79502a0c35 100644 --- a/docs/designs/ptoas-compiler-dso-wheel-linking.md +++ b/docs/designs/ptoas-compiler-dso-wheel-linking.md @@ -109,8 +109,7 @@ _core │ ├── ptoas.cpp │ ├── driver.cpp │ ├── VPTOHostStubEmission.cpp -│ ├── ObjectEmission.cpp -│ └── TilelangDaemon.cpp +│ └── ObjectEmission.cpp ├── PTOCAPI └── PTOASPythonCAPI ``` diff --git a/docs/designs/ptoas-python-launcher-layout.md b/docs/designs/ptoas-python-launcher-layout.md index 1a9d3d2d29..950ed1b070 100644 --- a/docs/designs/ptoas-python-launcher-layout.md +++ b/docs/designs/ptoas-python-launcher-layout.md @@ -139,7 +139,7 @@ assemble Python packages from unrelated build directories. The current archive is built against CPython 3.11 and requires a CPython 3.11 interpreter. `bin/ptoas` adds the archive root to `sys.path`, then uses the same `ptoas._cli -> ptoas._core` path as the install tree. The packaged `ptodsl/` -tree supports the compiler's default PTODSL TileLib backend; it does not turn +tree supports the compiler's PTODSL TileLib implementation; it does not turn the archive into a normal pip-installable PTODSL distribution. Linux archives use package-relative and archive-relative `$ORIGIN` RPATHs; @@ -155,8 +155,10 @@ declared installation layout. CTest and direct developer-tree runs must set an explicit matching `PYTHONPATH`; PTODSL does not guess repository, LLVM build, or PTOAS install paths at import time. -TileOp expansion remains a lazy, separate daemon process. The PTOAS CLI passes -the packaged PTODSL root and the active Python executable to the native driver, -which starts the daemon only when expansion is required. Keeping the daemon out -of the compiler process also prevents independently packaged MLIR/LLVM Python -bindings from registering runtime state in the native PTOAS process. +TileOp expansion runs in the CLI's existing Python process. `_core.main` creates +one Python-owned MLIR context for the compilation session, and the native driver +borrows that exact context. The in-process TileLib service materializes a source +module in the shared context and clones it into native ownership before the +Python module owner can be released. The packaged MLIR bindings and PTOAS must +therefore remain one ABI-matched runtime rather than independently replaceable +components. diff --git a/docs/designs/ptodsl-tilelib-template-selection-design.md b/docs/designs/ptodsl-tilelib-template-selection-design.md index 16e8b7f503..f5055c1e53 100644 --- a/docs/designs/ptodsl-tilelib-template-selection-design.md +++ b/docs/designs/ptodsl-tilelib-template-selection-design.md @@ -2,10 +2,7 @@ ## Background -PTOAS currently supports two TileLib backends for VPTO tile-op expansion: - -- `tilelang`, the legacy TileLangDSL template implementation. -- `ptodsl`, the PTODSL-native template implementation. +PTOAS uses the PTODSL-native TileLib implementation for VPTO tile-op expansion. A tile op may have several legal implementations for the same op name. Those implementations can differ by dtype, layout, memory space, @@ -42,7 +39,7 @@ in the ISA and user guide documents, not here. ## Pipeline -The PTODSL TileLib path has two compiler interactions with the Python daemon. +The PTODSL TileLib path has two interactions with the in-process Python service. ```text TileOp in MLIR @@ -50,7 +47,7 @@ TileOp in MLIR | InsertTemplateAttributes | - reconstruct operand specs from MLIR | - collect context attributes - | - ask the PTODSL daemon for legal candidates + | - ask the PTODSL service for legal candidates | - store compact candidate metadata on the TileOp v TileOp with candidates attr @@ -58,8 +55,8 @@ TileOp with candidates attr | ExpandTileOp | - build a specialization key from current MLIR operands and attrs | - choose candidate 0 from the compact candidates attr - | - ask the daemon to render that candidate - | - clone the generated helper and replace the TileOp with func.call + | - ask the service to materialize that candidate in the shared context + | - import the generated entry/helpers and replace the TileOp with func.call v VPTO-facing IR ``` @@ -102,7 +99,7 @@ remain in Python metadata for selection, diagnostics, and future tooling. ## Operand Specs Both `InsertTemplateAttributes` and `ExpandTileOp` reconstruct operand specs -from MLIR. The JSON shape sent to the daemon is deliberately close to +from MLIR. The JSON shape sent to the Python service is deliberately close to `TileSpec`, `ViewSpec`, `ScalarSpec`, and `VectorSpec`. | Operand kind | Required metadata | @@ -137,7 +134,7 @@ template is considered ported. ## Candidate Legality And Ranking -The daemon loads only the template module for the requested op and target. It +The service loads only the template module for the requested op and target. It then evaluates each registered candidate: 1. Bind positional MLIR operands to the template parameter names. @@ -149,7 +146,7 @@ then evaluates each registered candidate: 7. Run custom constraint predicates. 8. Sort legal candidates by descending priority. -If no candidate is legal, the daemon reports a `NoMatchingTemplate` error with +If no candidate is legal, the service reports a `NoMatchingTemplate` error with per-candidate reasons. If multiple candidates tie for the highest priority and no explicit candidate is requested, the registry reports ambiguity rather than silently picking one. @@ -170,7 +167,7 @@ TileOp. Each entry contains: - `tail` This attribute is intentionally not a copy of the full Python metadata object. -Legality has already happened in the daemon. The IR only needs a stable list of +Legality has already happened in the service. The IR only needs a stable list of legal render targets and the small amount of metadata consumed by downstream passes. @@ -180,8 +177,9 @@ Python metadata. Add a field only when a C++ pass or IR-level test consumes it. ## Expansion And Specialization `ExpandTileOp` uses the first candidate in the compact candidate list. For -PTODSL, it passes the selected candidate name back to the daemon so rendering -cannot accidentally choose a different legal template after the metadata pass. +PTODSL, it passes the selected candidate name back to the service so +materialization cannot accidentally choose a different legal template after the +metadata pass. The specialization key deduplicates generated helpers inside one module. It must include every input that can change the rendered helper body: diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index c1d9e82a14..50d9fdfd20 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -23,6 +23,7 @@ #include "llvm/ADT/StringRef.h" #include "mlir/Pass/Pass.h" #include "PTO/IR/PTODialect.h" +#include "PTO/Transforms/TileLibService.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -134,9 +135,10 @@ std::unique_ptr createVMINormalizeSignlessIntToUnsignedPass(); std::unique_ptr createVMIToVPTOPass(); std::unique_ptr createInsertTemplateAttributesPass(); std::unique_ptr createInsertTemplateAttributesPass( - const InsertTemplateAttributesOptions &options); + std::shared_ptr tileLibService); std::unique_ptr createExpandTileOpPass(); -std::unique_ptr createExpandTileOpPass(const ExpandTileOpOptions &options); +std::unique_ptr +createExpandTileOpPass(std::shared_ptr tileLibService); std::unique_ptr createFoldTileBufIntrinsicsPass(); std::unique_ptr createFoldTileBufIntrinsicsPass(llvm::StringRef foldMode); std::unique_ptr createPTOCanonicalizeIRPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 7b08319dda..a21ec988e2 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -533,40 +533,26 @@ def InsertTemplateAttributes : Pass<"pto-insert-template-attributes", "ModuleOp"> { let summary = "Attach legal PTODSL template candidates to tile operations"; let description = [{ - Queries the PTODSL TileLib daemon for legal template candidates and stores - the compact candidate list on each tile operation as the `candidates` - attribute. Each candidate contains only id, name, loop_depth, postupdate, - and tail metadata. + Queries the compiler's in-process PTODSL TileLib service for legal template + candidates and stores the compact candidate list on each tile operation as + the `candidates` attribute. Each candidate contains only id, name, + loop_depth, postupdate, and tail metadata. }]; let constructor = "mlir::pto::createInsertTemplateAttributesPass()"; let dependentDialects = [ "mlir::pto::PTODialect", "mlir::func::FuncDialect" ]; - let options = [ - Option<"pythonExe", "python-exe", "std::string", - /*default=*/"\"python3\"", - "Python executable for TileLib metadata invocation">, - Option<"daemonSocketPath", "daemon-socket-path", "std::string", - /*default=*/"\"\"", - "Path to the PTODSL TileLib daemon Unix socket">, - Option<"tileLibPkgPath", "tile-lib-pkg-path", "std::string", - /*default=*/"\"\"", - "PYTHONPATH root for PTODSL">, - Option<"daemonHelperModule", "daemon-helper-module", "std::string", - /*default=*/"\"ptodsl.tilelib.serving.helper\"", - "Python module used for daemon metadata RPC calls"> - ]; } def ExpandTileOp : Pass<"pto-expand-tile-op", "ModuleOp"> { let summary = "Expand tile ops into calls to TileLib template functions"; let description = [{ - Expands tile-level operations (pto.tadd, pto.tsub, etc.) by invoking the - selected Python TileLib backend to instantiate template libraries. The - generated template functions use tile_buf parameters and contain - vector-level implementations (pto.vecscope, pto.vlds, pto.vadd, - pto.vsts, etc.). + Expands tile-level operations (pto.tadd, pto.tsub, etc.) by asking the + compiler's in-process PTODSL TileLib service to instantiate template + libraries in the current MLIRContext. The generated template functions use + tile_buf parameters and contain vector-level implementations (pto.vecscope, + pto.vlds, pto.vadd, pto.vsts, etc.). Each tile op is replaced by a func.call to the generated template function, with tile_buf operands passed directly (no type bridging). @@ -584,29 +570,12 @@ def ExpandTileOp : Pass<"pto-expand-tile-op", "ModuleOp"> { "mlir::scf::SCFDialect", "mlir::vector::VectorDialect" ]; - let options = [ - Option<"pythonExe", "python-exe", "std::string", - /*default=*/"\"python3\"", - "Python executable for TileLib invocation">, - Option<"daemonSocketPath", "daemon-socket-path", "std::string", - /*default=*/"\"\"", - "Path to Unix domain socket for daemon RPC">, - Option<"tileLibBackend", "tile-lib-backend", "std::string", - /*default=*/"\"ptodsl\"", - "TileLib backend: ptodsl">, - Option<"tileLibPkgPath", "tile-lib-pkg-path", "std::string", - /*default=*/"\"\"", - "PYTHONPATH root for the selected TileLib backend">, - Option<"daemonHelperModule", "daemon-helper-module", "std::string", - /*default=*/"\"ptodsl.tilelib.serving.helper\"", - "Python module used for daemon helper RPC calls"> - ]; } def FoldTileBufIntrinsics : Pass<"pto-fold-tile-buf-intrinsics", "mlir::func::FuncOp"> { let summary = "Fold structured-view intrinsics after template inlining"; let description = [{ - After TileLang DSL template functions are inlined, the IR contains + After PTODSL template functions are inlined, the IR contains structured-view intrinsics whose operands are now bound to concrete values. This pass resolves them: diff --git a/include/PTO/Transforms/TileLibService.h b/include/PTO/Transforms/TileLibService.h new file mode 100644 index 0000000000..c28b19654a --- /dev/null +++ b/include/PTO/Transforms/TileLibService.h @@ -0,0 +1,54 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef MLIR_DIALECT_PTO_TRANSFORMS_TILELIBSERVICE_H +#define MLIR_DIALECT_PTO_TRANSFORMS_TILELIBSERVICE_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Support/LogicalResult.h" + +#include +#include + +namespace mlir::pto { + +/// Pure-data request used by the in-process TileLib materializer. The JSON +/// fields are request data only; generated MLIR never crosses this interface as +/// text. Keeping this interface independent of pybind11 allows transform passes +/// to remain usable from native tests and non-Python hosts. +struct TileLibMaterializationRequest { + std::string target; + std::string op; + std::string operandSpecsJson; + std::string contextAttrsJson; + std::string candidateId; +}; + +struct TileLibMaterialization { + OwningOpRef module; + std::string entrySymbol; +}; + +/// C++ ownership boundary for a TileLib implementation. Implementations return +/// a C++-owned source ModuleOp in the requested context. The caller may then +/// clone/import its generated functions into the caller module. +class TileLibService { +public: + virtual ~TileLibService() = default; + + virtual FailureOr + getMetadata(const TileLibMaterializationRequest &request) = 0; + + virtual FailureOr + materialize(const TileLibMaterializationRequest &request, + MLIRContext &context) = 0; +}; + +} // namespace mlir::pto + +#endif // MLIR_DIALECT_PTO_TRANSFORMS_TILELIBSERVICE_H diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 471afcb928..a67acd9599 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -9,8 +9,8 @@ //===- ExpandTileOp.cpp ---------------------------------------------------===// //===----------------------------------------------------------------------===// // -// Expand tile-level ops (pto.tadd, pto.tsub, ...) by invoking the selected -// Python TileLib backend to instantiate template libraries. +// Expand tile-level ops (pto.tadd, pto.tsub, ...) by materializing PTODSL +// template libraries in the compiler's host Python interpreter. // // The generated template functions use tile_buf parameters. After this pass, // the Inline pass inlines the template body, and FoldTileBufIntrinsics @@ -20,18 +20,18 @@ // 1. Extract SpecKey from ALL operands' tile_buf types. // 2. For PTODSL, read candidates attached by InsertTemplateAttributes and // select the first candidate still present. -// 3. Invoke the selected TileLib helper to generate a specialized MLIR -// function (with tile_buf parameters). -// 4. Parse the generated MLIR and clone the function into the module. +// 3. Ask the in-process TileLib service to build a source module in the same +// MLIRContext. +// 4. Clone its entry/helper functions into the caller module. // 5. Replace the original tile op with func.call, passing tile_buf // operands directly (no type bridging needed). // #include "PTO/IR/PTO.h" #include "PTO/IR/PTOTypeUtils.h" -#include "PTO/Support/PythonExecutable.h" #include "PTO/Transforms/Passes.h" #include "PTO/Transforms/TileOpExpansionUtils.h" +#include "PTO/Transforms/TileLibService.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" @@ -44,7 +44,6 @@ #include "mlir/IR/IRMapping.h" #include "mlir/IR/SymbolTable.h" #include "mlir/Pass/Pass.h" -#include "mlir/Parser/Parser.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallVector.h" @@ -53,20 +52,10 @@ #include "llvm/ADT/StringSet.h" #include "llvm/ADT/StringExtras.h" #include "llvm/ADT/StringRef.h" -#include "llvm/Support/FileSystem.h" -#include "llvm/Support/MemoryBuffer.h" -#include "llvm/Support/Path.h" -#include "llvm/Support/Program.h" #include "llvm/Support/raw_ostream.h" -#include #include #include -#include - -extern "C" { -extern char **environ; -} using namespace mlir; @@ -792,19 +781,13 @@ static std::optional buildSpecKey(Operation *op) { // ExpandState: runtime state for a single pass invocation. // ============================================================================ struct ExpandState { - std::vector> parsedModules; // Keep parsed modules alive + std::shared_ptr tileLibService; - std::string tileLibPkgPath; - std::string daemonHelperModule; - std::string pythonExe; - std::string daemonSocketPath; - - std::optional - invokeTileLibHelper(const SpecKey &key, StringRef candidateId = {}); func::FuncOp invokeTileLib(const SpecKey &key, Operation *tileOp, ModuleOp mod, MLIRContext *ctx); - func::FuncOp invokeTileLibDaemon(const SpecKey &key, StringRef candidateId, - ModuleOp mod, MLIRContext *ctx); + func::FuncOp invokeInProcessTileLib(const SpecKey &key, + StringRef candidateId, ModuleOp mod, + MLIRContext *ctx); LogicalResult expandTileOpsInFunction(func::FuncOp func, ModuleOp mod, MLIRContext *ctx); @@ -817,7 +800,17 @@ struct ExpandTileOpPass : public mlir::pto::impl::ExpandTileOpBase { using ExpandTileOpBase::ExpandTileOpBase; + explicit ExpandTileOpPass( + std::shared_ptr tileLibService) + : tileLibService(std::move(tileLibService)) {} + + ExpandTileOpPass(const ExpandTileOpPass &other) + : ExpandTileOpBase(other), + tileLibService(other.tileLibService) {} + void runOnOperation() override; + + std::shared_ptr tileLibService; }; /// Serialize a JSON array of integers. @@ -981,185 +974,116 @@ static std::string buildContextAttrsJson(const SpecKey &key) { } // ============================================================================ -// Invoke the configured one-shot helper and return its stdout. +// Materialize PTODSL in the host Python interpreter and import its functions. +// The service clones the Python-owned source module before returning, so this +// pass only handles C++-owned IR in the current MLIRContext. // ============================================================================ -std::optional -ExpandState::invokeTileLibHelper(const SpecKey &key, - StringRef candidateId) { - auto pythonPath = pto::resolvePythonExecutable(pythonExe); - if (!pythonPath) { - llvm::errs() << "ExpandTileOp: cannot find '" << pythonExe << "'\n"; - return std::nullopt; - } - - std::string operandSpecsJson = buildOperandSpecsJson(key); - std::string contextAttrsJson = buildContextAttrsJson(key); - if (key.targetArch.empty()) { - llvm::errs() << "ExpandTileOp: missing pto.target_arch module attribute\n"; - return std::nullopt; - } - - SmallString<128> tmpPath; - int tmpFD; - if (auto ec = llvm::sys::fs::createTemporaryFile("tilelib_helper", "out", - tmpFD, tmpPath)) { - llvm::errs() << "ExpandTileOp: cannot create temp file: " - << ec.message() << "\n"; - return std::nullopt; - } - ::close(tmpFD); - - std::string opName = "pto." + key.opName; - SmallVector args = { - *pythonPath, "-m", daemonHelperModule, - "--socket", daemonSocketPath, - "--target", key.targetArch, - "--op", opName, - "--operand-specs", operandSpecsJson, - }; - if (!key.contextAttrs.empty()) { - args.push_back("--context-attrs"); - args.push_back(contextAttrsJson); - } - if (!candidateId.empty()) { - args.push_back("--candidate-id"); - args.push_back(candidateId); - } - - std::optional redirects[] = {std::nullopt, StringRef(tmpPath), - std::nullopt}; - - SmallVector envp; - std::string pythonPathEnv; - std::vector envStorage; - bool hasPythonPath = !tileLibPkgPath.empty(); - if (hasPythonPath) { - const char *existingPath = ::getenv("PYTHONPATH"); - pythonPathEnv = "PYTHONPATH=" + tileLibPkgPath; - if (existingPath && existingPath[0] != '\0') { - pythonPathEnv += ":"; - pythonPathEnv += existingPath; - } - for (char **e = environ; *e; ++e) { - StringRef entry(*e); - if (entry.starts_with("PYTHONPATH=")) - continue; - envStorage.push_back(std::string(entry)); - } - envStorage.push_back(pythonPathEnv); - for (auto &s : envStorage) - envp.push_back(s); - } - - std::string errMsg; - int rc = llvm::sys::ExecuteAndWait( - *pythonPath, args, - hasPythonPath ? std::optional>(envp) : std::nullopt, - redirects, /*secondsToWait=*/30, /*memoryLimit=*/0, &errMsg); - - if (rc != 0) { - llvm::errs() << "ExpandTileOp: daemon helper instantiate failed (rc=" - << rc - << "): " << errMsg << "\n"; - llvm::sys::fs::remove(tmpPath); - return std::nullopt; - } +func::FuncOp ExpandState::invokeInProcessTileLib(const SpecKey &key, + StringRef candidateId, + ModuleOp mod, + MLIRContext *ctx) { + if (!tileLibService) + return nullptr; - auto bufOrErr = llvm::MemoryBuffer::getFile(tmpPath); - llvm::sys::fs::remove(tmpPath); - if (!bufOrErr) { - llvm::errs() << "ExpandTileOp: cannot read daemon output\n"; - return std::nullopt; - } - std::string output = (*bufOrErr)->getBuffer().str(); - if (output.empty()) { - llvm::errs() << "ExpandTileOp: empty daemon output\n"; - return std::nullopt; + pto::TileLibMaterializationRequest request; + request.target = key.targetArch; + request.op = "pto." + key.opName; + request.operandSpecsJson = buildOperandSpecsJson(key); + request.contextAttrsJson = buildContextAttrsJson(key); + request.candidateId = candidateId.str(); + + FailureOr materializationOr = + tileLibService->materialize(request, *ctx); + if (failed(materializationOr)) { + llvm::errs() << "ExpandTileOp: in-process PTODSL materialization failed\n"; + return nullptr; } - return output; -} -// ============================================================================ -// Invoke the daemon RPC to generate a specialized template function. -// ============================================================================ -func::FuncOp ExpandState::invokeTileLibDaemon(const SpecKey &key, - StringRef candidateId, - ModuleOp mod, - MLIRContext *ctx) { - auto mlirText = invokeTileLibHelper(key, candidateId); - if (!mlirText) + pto::TileLibMaterialization materialization = + std::move(*materializationOr); + OwningOpRef sourceModule = std::move(materialization.module); + if (!sourceModule || sourceModule->getContext() != ctx) { + llvm::errs() << "ExpandTileOp: in-process PTODSL returned a module from " + "a different MLIRContext\n"; return nullptr; + } - // Parse the rendered MLIR. - auto parsedMod = parseSourceString(*mlirText, ctx); - if (!parsedMod) { - llvm::errs() << "ExpandTileOp: failed to parse daemon output\n"; + auto sourceEntry = sourceModule->lookupSymbol( + materialization.entrySymbol); + if (!sourceEntry) { + llvm::errs() << "ExpandTileOp: in-process PTODSL entry symbol @" + << materialization.entrySymbol << " was not found\n"; return nullptr; } - // 9. Clone the generated function set into the target module. - auto parsedFuncs = parsedMod->getOps(); - if (parsedFuncs.empty()) { - llvm::errs() << "ExpandTileOp: no func.func in daemon output\n"; + SmallVector sourceFuncs; + for (func::FuncOp fn : sourceModule->getOps()) + sourceFuncs.push_back(fn); + if (sourceFuncs.empty()) { + llvm::errs() << "ExpandTileOp: in-process PTODSL returned no func.func\n"; return nullptr; } - // Create builder and set insertion point to insert functions into module - OpBuilder builder(ctx); - builder.setInsertionPointToEnd(mod.getBody()); - - llvm::StringMap renamedSymbols; - SmallVector clonedFuncs; - std::string uniqueName = buildUniqueFunctionBaseName(key); if (!candidateId.empty()) uniqueName += "__" + candidateId.str(); + SymbolTable targetSymTable(mod); if (auto existingFunc = targetSymTable.lookup(uniqueName)) return cast(existingFunc); - for (auto [index, fn] : llvm::enumerate(parsedFuncs)) { - // Use builder.clone() to insert into module body + llvm::StringMap plannedSymbols; + for (func::FuncOp fn : sourceFuncs) { + std::string newName = fn == sourceEntry + ? uniqueName + : uniqueName + "__" + std::string(fn.getSymName()); + if (targetSymTable.lookup(newName)) { + llvm::errs() << "ExpandTileOp: imported PTODSL symbol collision at @" + << newName << "\n"; + return nullptr; + } + plannedSymbols[fn.getSymName()] = std::move(newName); + } + + OpBuilder builder(ctx); + builder.setInsertionPointToEnd(mod.getBody()); + SmallVector clonedFuncs; + for (func::FuncOp fn : sourceFuncs) { IRMapping mapping; auto cloned = cast(builder.clone(*fn, mapping)); - std::string newName; - if (index == 0) { - newName = uniqueName; - } else { - newName = uniqueName + "__" + std::string(fn.getSymName()); - } - renamedSymbols[fn.getSymName()] = newName; - cloned.setName(newName); - - // Set visibility to Private for template functions (required for inline pass) + cloned.setName(plannedSymbols.lookup(fn.getSymName())); cloned.setVisibility(SymbolTable::Visibility::Private); - clonedFuncs.push_back(cloned); } for (func::FuncOp fn : clonedFuncs) { - fn.walk([&](func::CallOp call) { - StringRef callee = call.getCallee(); - if (callee.empty()) - return; - auto renameIt = renamedSymbols.find(callee); - if (renameIt == renamedSymbols.end()) - return; - call.setCallee(renameIt->second); - }); + for (const auto &renamed : plannedSymbols) { + if (failed(SymbolTable::replaceAllSymbolUses( + StringAttr::get(ctx, renamed.getKey()), + StringAttr::get(ctx, renamed.getValue()), fn))) { + llvm::errs() << "ExpandTileOp: failed to rewrite imported symbol @" + << renamed.getKey() << " in @" << fn.getSymName() + << "\n"; + for (func::FuncOp imported : clonedFuncs) + imported.erase(); + return nullptr; + } + } } - auto cloned = clonedFuncs.front(); - if (!cloned->hasAttr("pto.tilelang.instance")) { - llvm::errs() << "ExpandTileOp: warning: daemon output function @" - << cloned.getSymName() + func::FuncOp entry = mod.lookupSymbol(uniqueName); + if (!entry) { + llvm::errs() << "ExpandTileOp: failed to import PTODSL entry @" + << materialization.entrySymbol << "\n"; + return nullptr; + } + if (!entry->hasAttr("pto.tilelang.instance")) { + llvm::errs() << "ExpandTileOp: warning: in-process PTODSL entry @" + << entry.getSymName() << " missing pto.tilelang.instance attribute\n"; } - // Keep the parsed module alive. - parsedModules.push_back(std::move(parsedMod)); - - return cloned; + return entry; } // ============================================================================ @@ -1168,8 +1092,9 @@ func::FuncOp ExpandState::invokeTileLibDaemon(const SpecKey &key, func::FuncOp ExpandState::invokeTileLib(const SpecKey &key, Operation *tileOp, ModuleOp mod, MLIRContext *ctx) { - if (daemonSocketPath.empty()) { - llvm::errs() << "ExpandTileOp: PTODSL backend requires its daemon\n"; + if (!tileLibService) { + tileOp->emitError( + "ExpandTileOp PTODSL backend requires an in-process service"); return nullptr; } @@ -1190,13 +1115,7 @@ func::FuncOp ExpandState::invokeTileLib(const SpecKey &key, return nullptr; } - func::FuncOp daemonResult = - invokeTileLibDaemon(key, selectedName.getValue(), mod, ctx); - if (daemonResult) - return daemonResult; - - llvm::errs() << "ExpandTileOp: PTODSL daemon RPC failed\n"; - return nullptr; + return invokeInProcessTileLib(key, selectedName.getValue(), mod, ctx); } // ============================================================================ @@ -1222,7 +1141,7 @@ LogicalResult ExpandState::expandTileOpsInFunction(func::FuncOp func, return failure(); } - // Invoke the selected TileLib backend (with daemon-side caching). + // Materialize the selected PTODSL template in-process. func::FuncOp dslFn = invokeTileLib(*specKeyOpt, op, mod, ctx); if (!dslFn) { StringRef opName = getTileOpName(op); @@ -1271,24 +1190,14 @@ void ExpandTileOpPass::runOnOperation() { if (!hasExpandableOps) return; - if (tileLibBackend != "ptodsl") { - mod.emitError("ExpandTileOp received unsupported tile-lib-backend '" + - std::string(tileLibBackend) + "'"); - signalPassFailure(); - return; - } - - if (daemonSocketPath.empty()) { - mod.emitError("ExpandTileOp requires a running PTODSL TileLib daemon"); + if (!tileLibService) { + mod.emitError("ExpandTileOp PTODSL backend requires an in-process service"); signalPassFailure(); return; } ExpandState state; - state.tileLibPkgPath = std::string(tileLibPkgPath); - state.daemonHelperModule = std::string(daemonHelperModule); - state.pythonExe = std::string(pythonExe); - state.daemonSocketPath = std::string(daemonSocketPath); + state.tileLibService = tileLibService; for (auto func : mod.getOps()) { if (func.isExternal()) @@ -1307,9 +1216,9 @@ std::unique_ptr createExpandTileOpPass() { return std::make_unique(); } -std::unique_ptr -createExpandTileOpPass(const ExpandTileOpOptions &options) { - return std::make_unique(options); +std::unique_ptr createExpandTileOpPass( + std::shared_ptr tileLibService) { + return std::make_unique(std::move(tileLibService)); } } // namespace pto diff --git a/lib/PTO/Transforms/InsertTemplateAttributes.cpp b/lib/PTO/Transforms/InsertTemplateAttributes.cpp index f1069d4f5b..cefd2aeb2f 100644 --- a/lib/PTO/Transforms/InsertTemplateAttributes.cpp +++ b/lib/PTO/Transforms/InsertTemplateAttributes.cpp @@ -8,7 +8,6 @@ #include "PTO/IR/PTO.h" #include "PTO/IR/PTOTypeUtils.h" -#include "PTO/Support/PythonExecutable.h" #include "PTO/Transforms/Passes.h" #include "PTO/Transforms/TileOpExpansionUtils.h" @@ -21,27 +20,17 @@ #include "mlir/Pass/Pass.h" #include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallString.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringExtras.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/Error.h" -#include "llvm/Support/FileSystem.h" #include "llvm/Support/JSON.h" -#include "llvm/Support/MemoryBuffer.h" -#include "llvm/Support/Program.h" #include "llvm/Support/raw_ostream.h" -#include #include #include -#include #include -extern "C" { -extern char **environ; -} - using namespace mlir; namespace mlir { @@ -767,121 +756,6 @@ getTargetArch(Operation *operation) { return std::nullopt; } -static std::optional -invokeMetadataHelper(Operation *operation, StringRef pythonExe, - StringRef daemonSocketPath, StringRef tileLibPkgPath, - StringRef daemonHelperModule) { - auto pythonPath = pto::resolvePythonExecutable(pythonExe); - if (!pythonPath) { - operation->emitError("InsertTemplateAttributes cannot find Python '") - << pythonExe << "'"; - return std::nullopt; - } - - auto target = getTargetArch(operation); - auto operandSpecs = buildOperandSpecsJson(operation); - if (!target || !operandSpecs) - return std::nullopt; - std::string contextAttrs = buildContextAttrsJson(operation); - - llvm::SmallString<128> outputPath; - int outputFd; - if (auto error = llvm::sys::fs::createTemporaryFile( - "tilelib_metadata", "json", outputFd, outputPath)) { - operation->emitError("InsertTemplateAttributes cannot create temporary " - "metadata output: ") - << error.message(); - return std::nullopt; - } - ::close(outputFd); - - llvm::SmallString<128> errorPath; - int errorFd; - if (auto error = llvm::sys::fs::createTemporaryFile( - "tilelib_metadata", "err", errorFd, errorPath)) { - llvm::sys::fs::remove(outputPath); - operation->emitError("InsertTemplateAttributes cannot create temporary " - "metadata error output: ") - << error.message(); - return std::nullopt; - } - ::close(errorFd); - - std::string opName = operation->getName().getStringRef().str(); - SmallVector args = { - *pythonPath, "-m", daemonHelperModule, - "--method", "get_metadata", "--socket", - daemonSocketPath, "--target", *target, - "--op", opName, "--operand-specs", - *operandSpecs, - }; - if (contextAttrs != "{}") { - args.push_back("--context-attrs"); - args.push_back(contextAttrs); - } - - std::optional redirects[] = { - std::nullopt, - StringRef(outputPath), - StringRef(errorPath), - }; - - SmallVector environment; - std::string pythonPathEnvironment; - std::vector environmentStorage; - bool hasPythonPath = !tileLibPkgPath.empty(); - if (hasPythonPath) { - const char *existingPath = ::getenv("PYTHONPATH"); - pythonPathEnvironment = "PYTHONPATH=" + tileLibPkgPath.str(); - if (existingPath && existingPath[0] != '\0') - pythonPathEnvironment += ":" + std::string(existingPath); - - for (char **entry = environ; *entry; ++entry) { - StringRef value(*entry); - if (!value.starts_with("PYTHONPATH=")) - environmentStorage.push_back(value.str()); - } - environmentStorage.push_back(pythonPathEnvironment); - for (std::string &value : environmentStorage) - environment.push_back(value); - } - - std::string errorMessage; - int result = llvm::sys::ExecuteAndWait( - *pythonPath, args, - hasPythonPath - ? std::optional>(environment) - : std::nullopt, - redirects, /*secondsToWait=*/30, /*memoryLimit=*/0, &errorMessage); - if (result != 0) { - auto errorOutput = llvm::MemoryBuffer::getFile(errorPath); - llvm::sys::fs::remove(outputPath); - llvm::sys::fs::remove(errorPath); - - std::string detail; - if (errorOutput) - detail = errorOutput.get()->getBuffer().trim().str(); - if (detail.empty()) - detail = errorMessage; - if (detail.empty()) - detail = "helper exited with status " + std::to_string(result); - - operation->emitError("InsertTemplateAttributes metadata RPC failed: ") - << detail; - return std::nullopt; - } - - auto output = llvm::MemoryBuffer::getFile(outputPath); - llvm::sys::fs::remove(outputPath); - llvm::sys::fs::remove(errorPath); - if (!output) { - operation->emitError( - "InsertTemplateAttributes cannot read metadata output"); - return std::nullopt; - } - return (*output)->getBuffer().str(); -} - static FailureOr parseCandidateAttributes(Operation *operation, StringRef metadataJson) { auto parsed = llvm::json::parse(metadataJson); @@ -981,6 +855,14 @@ struct InsertTemplateAttributesPass InsertTemplateAttributesPass> { using InsertTemplateAttributesBase::InsertTemplateAttributesBase; + explicit InsertTemplateAttributesPass( + std::shared_ptr tileLibService) + : tileLibService(std::move(tileLibService)) {} + + InsertTemplateAttributesPass(const InsertTemplateAttributesPass &other) + : InsertTemplateAttributesBase(other), + tileLibService(other.tileLibService) {} + void runOnOperation() override { ModuleOp module = getOperation(); @@ -991,18 +873,27 @@ struct InsertTemplateAttributesPass }); if (tileOperations.empty()) return; - if (daemonSocketPath.empty()) { + if (!tileLibService) { module.emitError( - "InsertTemplateAttributes requires a PTODSL daemon socket"); + "InsertTemplateAttributes requires an in-process PTODSL service"); return signalPassFailure(); } for (Operation *operation : tileOperations) { - auto metadata = invokeMetadataHelper( - operation, pythonExe, daemonSocketPath, tileLibPkgPath, - daemonHelperModule); - if (!metadata) + auto target = getTargetArch(operation); + auto operandSpecs = buildOperandSpecsJson(operation); + if (!target || !operandSpecs) return signalPassFailure(); + pto::TileLibMaterializationRequest request; + request.target = std::move(*target); + request.op = operation->getName().getStringRef().str(); + request.operandSpecsJson = std::move(*operandSpecs); + request.contextAttrsJson = buildContextAttrsJson(operation); + FailureOr metadata = tileLibService->getMetadata(request); + if (failed(metadata)) { + operation->emitError("in-process PTODSL metadata query failed"); + return signalPassFailure(); + } auto candidates = parseCandidateAttributes(operation, *metadata); if (failed(candidates)) @@ -1010,6 +901,8 @@ struct InsertTemplateAttributesPass operation->setAttr(kCandidatesAttr, *candidates); } } + + std::shared_ptr tileLibService; }; } // namespace @@ -1022,8 +915,9 @@ std::unique_ptr createInsertTemplateAttributesPass() { } std::unique_ptr createInsertTemplateAttributesPass( - const InsertTemplateAttributesOptions &options) { - return std::make_unique(options); + std::shared_ptr tileLibService) { + return std::make_unique( + std::move(tileLibService)); } } // namespace pto diff --git a/ptodsl/README.md b/ptodsl/README.md index 0dff856683..4303730f1e 100644 --- a/ptodsl/README.md +++ b/ptodsl/README.md @@ -86,18 +86,17 @@ LLVM_BUILD_DIR=/path/to/llvm/build ./quick_install.sh ## PTODSL TileLib backend -PTOAS uses the PTODSL TileLib daemon by default for VPTO tile-op expansion: +PTOAS uses its in-process PTODSL TileLib service for VPTO tile-op expansion: ```bash ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto \ input.pto -o - ``` -Wheel and CMake-tree launchers pass the Python root containing their own -installed or staged `ptodsl` package to the native driver. Use -`--ptodsl-pkg-path=/path/to/package/root` for an explicit command-line -override. PTODSL daemon failures are reported as errors and never fall back to -the TileLang implementation. +Wheel and CMake-tree launchers add their packaged `TileOps` resource directory +to the host interpreter while the native compiler runs. Template discovery, +metadata queries, and MLIR materialization therefore stay in one process and +do not require a Python executable, daemon socket, or package-path CLI option. `InsertTemplateAttributes` queries legal-candidate metadata before fusion and stores an ordered `candidates` array containing only `id`, `name`, diff --git a/ptodsl/docs/developer_guide/tilelib-debugging-playbook.md b/ptodsl/docs/developer_guide/tilelib-debugging-playbook.md index d5c4c015f8..8b5d7775c8 100644 --- a/ptodsl/docs/developer_guide/tilelib-debugging-playbook.md +++ b/ptodsl/docs/developer_guide/tilelib-debugging-playbook.md @@ -1,7 +1,7 @@ # PTODSL TileLib Debugging Playbook -This playbook is for diagnosing PTODSL TileLib failures while migrating from -TileLangDSL. It converts the migration scratch notes into a reusable workflow. +This playbook is for diagnosing PTODSL TileLib failures. It converts the +original migration scratch notes into a reusable workflow. The main rule: classify the failure before editing templates. A build failure, a candidate-selection failure, a tracing failure, and a wrong-output failure @@ -24,7 +24,6 @@ ninja -C build-llvm21 PTODSLPackage Run one smoke ST: ```bash -PTOAS_TILE_LIB_BACKEND=ptodsl \ python3 test/tilelang_st/script/run_all_st.py \ -r sim -v a5 \ -p build-llvm21/tools/ptoas/ptoas \ @@ -34,7 +33,6 @@ python3 test/tilelang_st/script/run_all_st.py \ Run one non-smoke ST: ```bash -PTOAS_TILE_LIB_BACKEND=ptodsl \ python3 test/tilelang_st/script/run_all_st.py \ -r sim -v a5 \ -p build-llvm21/tools/ptoas/ptoas \ @@ -44,7 +42,6 @@ python3 test/tilelang_st/script/run_all_st.py \ Run one named ST case when supported by `run_st.py`: ```bash -PTOAS_TILE_LIB_BACKEND=ptodsl \ python3 test/tilelang_st/script/run_st.py \ -r sim -v a5 \ -p build-llvm21/tools/ptoas/ptoas \ @@ -64,8 +61,7 @@ python3 test/tilelang_st/script/run_st.py \ | isolated case passes but full test fails | helper specialization cache or stale generated package | Do not assume every ST failure means the testcase is wrong. First check whether -the same case works with TileLangDSL and whether the PTODSL lowering has enough -metadata to reproduce the TileLangDSL behavior. +the PTODSL lowering has enough metadata to represent the testcase behavior. ## Candidate Selection Failures @@ -94,15 +90,15 @@ If legality appears correct but `ExpandTileOp` cannot expand: non-empty `candidates` attr. 2. Dump after passes that rewrite view/tile operands and confirm the attr is still attached. -3. Confirm candidate 0 has a `name` and that the daemon can render that - candidate directly. +3. Confirm candidate 0 has a `name` and that the in-process TileLib service can + materialize that candidate directly. Useful compiler-only command: ```bash build-llvm21/tools/ptoas/ptoas \ --pto-arch=a5 --pto-backend=vpto --emit-vpto \ - --tile-lib-backend=ptodsl --enable-insert-sync \ + --enable-insert-sync \ --mlir-print-ir-after=pto-expand-tile-op \ --mlir-print-ir-tree-dir=/tmp/_after_expand_ptodsl \ test/tilelang_st/npu/a5/src/st/testcase//.pto \ @@ -133,7 +129,7 @@ Compiler-only dump: ```bash build-llvm21/tools/ptoas/ptoas \ --pto-arch=a5 --pto-backend=vpto --emit-vpto \ - --tile-lib-backend=ptodsl --enable-insert-sync \ + --enable-insert-sync \ test/tilelang_st/npu/a5/src/st/testcase//.pto \ -o /tmp/_ptodsl.vpto ``` diff --git a/ptodsl/docs/developer_guide/tilelib-template-authoring.md b/ptodsl/docs/developer_guide/tilelib-template-authoring.md index df05cefd6c..2eca3b56e0 100644 --- a/ptodsl/docs/developer_guide/tilelib-template-authoring.md +++ b/ptodsl/docs/developer_guide/tilelib-template-authoring.md @@ -33,7 +33,7 @@ def template_tadd(src0, src1, dst): ... ``` -The function parameter order is the operand binding contract. The daemon binds +The function parameter order is the operand binding contract. The TileLib runtime binds MLIR operands positionally to these parameter names before evaluating constraints or rendering. If a TileLangDSL template had multiple callable forms, either match the ST operand order exactly or register separate PTODSL diff --git a/ptodsl/ptoas/_cli.py b/ptodsl/ptoas/_cli.py index a12158ad7d..90e5b5ecc0 100644 --- a/ptodsl/ptoas/_cli.py +++ b/ptodsl/ptoas/_cli.py @@ -31,13 +31,12 @@ def _load_native_module(): return _core -def _resolve_runtime_paths(native_module) -> tuple[Path, Path]: +def _resolve_tileops_dir(native_module) -> Path: module_file = getattr(native_module, "__file__", None) if not module_file: raise SystemExit("ptoas._core does not expose a module file") package_root = Path(module_file).resolve().parent - python_root = package_root.parent runtime_root = package_root / "_runtime" tileops_dir = runtime_root / "share" / "ptoas" / "TileOps" if not tileops_dir.is_dir(): @@ -45,32 +44,27 @@ def _resolve_runtime_paths(native_module) -> tuple[Path, Path]: "unable to locate packaged PTOAS TileOps resources: expected " f"{tileops_dir}" ) - return python_root, tileops_dir.resolve() - - -def _has_cli_option(arguments: Sequence[str], option: str) -> bool: - option_with_value = f"{option}=" - return any( - argument == option or argument.startswith(option_with_value) - for argument in arguments - ) + return tileops_dir.resolve() def launch(user_args: Sequence[str], *, wrapper: Path | None = None) -> int: native_module = _load_native_module() - python_root, tileops_dir = _resolve_runtime_paths(native_module) + tileops_dir = _resolve_tileops_dir(native_module) wrapper = wrapper.resolve() if wrapper is not None else _resolve_wrapper_path() os.environ["PTOAS_BIN"] = str(wrapper) - os.environ["PTOAS_PYTHON_EXE"] = sys.executable argv = [str(wrapper)] - if not _has_cli_option(user_args, "--ptodsl-pkg-path"): - argv.extend(["--ptodsl-pkg-path", str(python_root)]) - if not _has_cli_option(user_args, "--tileops-pkg-path"): - argv.extend(["--tileops-pkg-path", str(tileops_dir.parent)]) argv.extend(user_args) - return int(native_module.main(argv)) + tileops_python_root = str(tileops_dir.parent) + inserted_tileops_root = tileops_python_root not in sys.path + if inserted_tileops_root: + sys.path.insert(0, tileops_python_root) + try: + return int(native_module.main(argv)) + finally: + if inserted_tileops_root: + sys.path.remove(tileops_python_root) def main() -> int: diff --git a/ptodsl/ptodsl/tilelib/_compiler_runtime.py b/ptodsl/ptodsl/tilelib/_compiler_runtime.py new file mode 100644 index 0000000000..23a183273d --- /dev/null +++ b/ptodsl/ptodsl/tilelib/_compiler_runtime.py @@ -0,0 +1,72 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""In-process PTODSL TileLib materialization entry point.""" + +from __future__ import annotations + +import json + +from ._selection import _select_descriptor_and_specs, metadata_request + + +def metadata( + target: str, + op: str, + operand_specs_json: str, + context_attrs_json: str, +) -> str: + """Return candidate metadata JSON without any daemon transport.""" + try: + operand_specs = json.loads(operand_specs_json) + context_attrs = json.loads(context_attrs_json or "{}") + except json.JSONDecodeError as exc: + raise ValueError(f"invalid TileLib metadata request: {exc}") from exc + return json.dumps( + metadata_request(target, op, operand_specs, context_attrs), + separators=(",", ":"), + sort_keys=True, + ) + + +def materialize( + target: str, + op: str, + operand_specs_json: str, + context_attrs_json: str, + candidate_id: str | None, + context, +): + """Return ``(source_module, entry_symbol)`` in *context*. + + JSON is retained only for the compact pure-data specialization request. No + MLIR text is produced or parsed by this path. + """ + try: + operand_specs = json.loads(operand_specs_json) + context_attrs = json.loads(context_attrs_json or "{}") + except json.JSONDecodeError as exc: + raise ValueError(f"invalid TileLib materialization request: {exc}") from exc + + descriptor, tile_specs = _select_descriptor_and_specs( + target, + op, + operand_specs, + context_attrs, + candidate_id or None, + ) + artifact = descriptor.specialize( + context_attrs=context_attrs, + **tile_specs, + ) + module = artifact.materialize(context) + module.operation.verify() + return module, descriptor.name + + +__all__ = ["materialize", "metadata"] diff --git a/ptodsl/ptodsl/tilelib/_render_runtime.py b/ptodsl/ptodsl/tilelib/_render_runtime.py index d925c76ce0..b79386bd81 100644 --- a/ptodsl/ptodsl/tilelib/_render_runtime.py +++ b/ptodsl/ptodsl/tilelib/_render_runtime.py @@ -176,8 +176,8 @@ def trace_entry(self, *args): rewritten(*args) # Custom golden-shaped container: single module(target_arch) + func(instance, kernel_kind). - def build_module(self): - ctx = make_context() + def build_module(self, context=None): + ctx = context if context is not None else make_context() with ctx, Location.unknown(): arg_types = list(self.compute_argument_types()) module, ir_fn = self._create_instance_module(arg_types) diff --git a/ptodsl/ptodsl/tilelib/serving/daemon.py b/ptodsl/ptodsl/tilelib/_selection.py similarity index 58% rename from ptodsl/ptodsl/tilelib/serving/daemon.py rename to ptodsl/ptodsl/tilelib/_selection.py index f52c4c6dcf..d34ebe51f7 100644 --- a/ptodsl/ptodsl/tilelib/serving/daemon.py +++ b/ptodsl/ptodsl/tilelib/_selection.py @@ -5,43 +5,18 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""PTODSL TileLib daemon for the ExpandTileOp Unix-socket RPC contract. - -The daemon owns template discovery, selection, specialization, rendering, and an -in-memory instance cache. PTODSL templates are loaded from the Python package, so -the daemon does not scan or depend on an external template directory. - -Run it with: - - python3 -m ptodsl.tilelib.serving.daemon --socket -""" +"""PTODSL TileLib candidate discovery, validation, and selection.""" from __future__ import annotations -import argparse -import json -import os -import signal -import socketserver -import threading - -from .. import constraints as _constraints -from .. import registry as _registry -from ..metadata import ScalarSpec, ScalarType, TileSpec, VectorSpec, ViewSpec +from . import constraints as _constraints +from . import registry as _registry +from .metadata import ScalarSpec, ScalarType, TileSpec, VectorSpec, ViewSpec from TileOps import load_template -from .wire import recv_message, send_message - - -def _remove_socket_path(socket_path: str) -> None: - """Remove an existing socket entry, including a broken symlink.""" - try: - os.unlink(socket_path) - except FileNotFoundError: - pass def _build_tile_specs(descriptor, operand_specs: list) -> dict: - """Map positional daemon operands onto a template's parameter names.""" + """Map positional compiler operands onto a template's parameter names.""" if not isinstance(operand_specs, list): raise TypeError("operand_specs must be a list") if len(operand_specs) != len(descriptor.param_names): @@ -101,7 +76,7 @@ def _build_tile_specs(descriptor, operand_specs: list) -> dict: if kind != "tile": raise NotImplementedError( - "PTODSL TileLib daemon currently supports tile, scalar, view, " + "PTODSL TileLib currently supports tile, scalar, view, " f"and vector operands; " f"operand {index} ({name!r}) has kind {kind!r}" ) @@ -296,183 +271,4 @@ def metadata_request( } -def render_request( - target: str, - op: str, - operand_specs: list, - context_attrs: dict | None = None, - candidate_id: str | None = None, -) -> str: - """Select and render one PTODSL template as MLIR text.""" - descriptor, tile_specs = _select_descriptor_and_specs( - target, - op, - operand_specs, - context_attrs, - candidate_id, - ) - return descriptor.specialize( - context_attrs=context_attrs or {}, - **tile_specs, - ).mlir_text() - - -class TileLibDaemonServer(socketserver.UnixStreamServer): - """Sequential Unix-socket RPC server with an in-memory render cache.""" - - def __init__(self, socket_path: str, max_entries: int = 1000): - if max_entries <= 0: - raise ValueError("max_entries must be greater than zero") - super().__init__(socket_path, _Handler) - os.chmod(socket_path, 0o600) - self._cache: dict[str, str] = {} - self._max_entries = max_entries - self._stats = {"hits": 0, "misses": 0, "evictions": 0} - - @property - def stats(self) -> dict: - """Return a snapshot of cache counters for diagnostics and tests.""" - return dict(self._stats) - - def dispatch(self, request: dict) -> dict: - if not isinstance(request, dict): - return {"success": False, "error": "request must be a JSON object"} - - method = request.get("method") - params = request.get("params") or {} - if not isinstance(params, dict): - return {"success": False, "error": "request params must be a JSON object"} - - try: - if method == "instantiate": - result = self._instantiate(**params) - elif method == "get_metadata": - result = self._get_metadata(**params) - elif method == "ping": - result = "pong" - elif method == "get_stats": - result = self._get_stats() - elif method == "clear": - result = self._clear() - else: - return {"success": False, "error": f"unknown method {method!r}"} - return {"success": True, "result": result} - except Exception as exc: - return { - "success": False, - "error": f"{type(exc).__name__}: {exc}", - } - - def _get_metadata(self, target, op, operand_specs, context_attrs=None): - return metadata_request(target, op, operand_specs, context_attrs) - - def _get_stats(self): - requests = self._stats["hits"] + self._stats["misses"] - total_entries = len(self._cache) - return { - **self._stats, - "entries": total_entries, - "total_entries": total_entries, - "max_entries": self._max_entries, - "hit_rate": self._stats["hits"] / requests if requests else 0.0, - } - - def _clear(self): - self._cache.clear() - return {"cleared": True} - - def _instantiate( - self, - target, - op, - operand_specs, - context_attrs=None, - candidate_id=None, - ): - key = json.dumps( - { - "target": target, - "op": op, - "operand_specs": operand_specs, - "context_attrs": context_attrs, - "candidate_id": candidate_id, - }, - sort_keys=True, - separators=(",", ":"), - ) - - cached = self._cache.get(key) - if cached is not None: - self._stats["hits"] += 1 - return cached - self._stats["misses"] += 1 - - mlir_text = render_request( - target, - op, - operand_specs, - context_attrs, - candidate_id, - ) - - if len(self._cache) >= self._max_entries: - self._cache.pop(next(iter(self._cache))) - self._stats["evictions"] += 1 - self._cache[key] = mlir_text - return mlir_text - - -class _Handler(socketserver.BaseRequestHandler): - def handle(self): - try: - request = recv_message(self.request) - except (ConnectionError, UnicodeDecodeError, ValueError): - return - send_message(self.request, self.server.dispatch(request)) - - -def _parse_args(argv): - parser = argparse.ArgumentParser(prog="ptodsl.tilelib.serving.daemon") - parser.add_argument("--socket", required=True) - parser.add_argument( - "--template-dir", - default=None, - help="accepted during migration but ignored; PTODSL templates are in-package", - ) - parser.add_argument("--max-entries", type=int, default=1000) - parser.add_argument("--verbose", action="store_true") - return parser.parse_args(argv) - - -def main(argv=None): - args = _parse_args(argv) - - _remove_socket_path(args.socket) - - server = TileLibDaemonServer(args.socket, max_entries=args.max_entries) - stop = threading.Event() - - def _request_shutdown(*_): - stop.set() - - signal.signal(signal.SIGTERM, _request_shutdown) - signal.signal(signal.SIGINT, _request_shutdown) - - thread = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - if args.verbose: - print(f"PTODSL TileLib daemon listening on {args.socket}", flush=True) - - try: - stop.wait() - finally: - server.shutdown() - server.server_close() - _remove_socket_path(args.socket) - - -if __name__ == "__main__": - main() - - -__all__ = ["TileLibDaemonServer", "main", "metadata_request", "render_request"] +__all__ = ["metadata_request"] diff --git a/ptodsl/ptodsl/tilelib/decorator.py b/ptodsl/ptodsl/tilelib/decorator.py index ba0f3db095..0cf4294353 100644 --- a/ptodsl/ptodsl/tilelib/decorator.py +++ b/ptodsl/ptodsl/tilelib/decorator.py @@ -56,6 +56,19 @@ def __init__(self, descriptor: TileTemplate, tile_specs: dict, context_attrs=Non self.tile_specs = tile_specs self.context_attrs = dict(context_attrs or {}) + def materialize(self, context): + """Build a fresh source module in the caller-provided context. + + The returned module remains Python-owned. Native callers must keep this + object alive until they have cloned/imported its generated functions. + This bypasses the context-bound ModuleArtifact cache. + """ + return _TemplateTrace( + self.descriptor, + self.tile_specs, + context_attrs=self.context_attrs, + ).build_module(context=context) + def tile_template(*, op, target="a5", name=None, dtypes=(), layouts=(), memory_spaces=(), constraints=(), priority=0, fusible=False, diff --git a/ptodsl/ptodsl/tilelib/serving/__init__.py b/ptodsl/ptodsl/tilelib/serving/__init__.py deleted file mode 100644 index 3ac3c8f6a3..0000000000 --- a/ptodsl/ptodsl/tilelib/serving/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software, you can redistribute it and/or modify it under the terms and conditions of -# CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -# See LICENSE in the root of the software repository for the full text of the License. -"""Unix-socket serving layer for the PTODSL TileLib.""" - -from .client import DaemonClient, DaemonError - - -def __getattr__(name): - # Keep daemon.py unloaded when executing it with ``python -m``. - if name in {"TileLibDaemonServer", "metadata_request", "render_request"}: - from .daemon import TileLibDaemonServer, metadata_request, render_request - - exports = { - "TileLibDaemonServer": TileLibDaemonServer, - "metadata_request": metadata_request, - "render_request": render_request, - } - return exports[name] - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -__all__ = [ - "DaemonClient", - "DaemonError", - "TileLibDaemonServer", - "metadata_request", - "render_request", -] diff --git a/ptodsl/ptodsl/tilelib/serving/client.py b/ptodsl/ptodsl/tilelib/serving/client.py deleted file mode 100644 index 9a4db37b0e..0000000000 --- a/ptodsl/ptodsl/tilelib/serving/client.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software, you can redistribute it and/or modify it under the terms and conditions of -# CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -# See LICENSE in the root of the software repository for the full text of the License. -"""Synchronous client for the PTODSL TileLib daemon.""" - -from __future__ import annotations - -import socket - -from .wire import recv_message, send_message - - -class DaemonError(Exception): - """An RPC reached the daemon but the requested operation failed.""" - - -class DaemonClient: - """Issue one daemon RPC per Unix-socket connection.""" - - def __init__(self, socket_path: str): - self.socket_path = socket_path - - def _call(self, method: str, params: dict | None = None): - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: - sock.connect(self.socket_path) - send_message(sock, {"method": method, "params": params or {}}) - response = recv_message(sock) - - if not response.get("success"): - raise DaemonError(response.get("error", "unknown daemon error")) - return response["result"] - - def ping(self): - return self._call("ping") - - def get_metadata(self, target, op, operand_specs, context_attrs=None): - return self._call( - "get_metadata", - { - "target": target, - "op": op, - "operand_specs": operand_specs, - "context_attrs": context_attrs or {}, - }, - ) - - def instantiate( - self, - target, - op, - operand_specs, - context_attrs=None, - candidate_id=None, - ): - return self._call( - "instantiate", - { - "target": target, - "op": op, - "operand_specs": operand_specs, - "context_attrs": context_attrs or {}, - "candidate_id": candidate_id, - }, - ) - - def get_stats(self): - return self._call("get_stats") - - def clear(self): - return self._call("clear") - - -__all__ = ["DaemonClient", "DaemonError"] diff --git a/ptodsl/ptodsl/tilelib/serving/helper.py b/ptodsl/ptodsl/tilelib/serving/helper.py deleted file mode 100644 index 5925556ef7..0000000000 --- a/ptodsl/ptodsl/tilelib/serving/helper.py +++ /dev/null @@ -1,73 +0,0 @@ -# Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software, you can redistribute it and/or modify it under the terms and conditions of -# CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -# See LICENSE in the root of the software repository for the full text of the License. -"""One-shot command-line client for the ExpandTileOp daemon contract. - -Example: - - python3 -m ptodsl.tilelib.serving.helper --socket --target a5 \ - --op pto.tadd --operand-specs '[...]' -""" - -from __future__ import annotations - -import argparse -import json -import sys - -from .client import DaemonClient, DaemonError - - -def main(argv=None): - parser = argparse.ArgumentParser(prog="ptodsl.tilelib.serving.helper") - parser.add_argument("--socket", required=True) - parser.add_argument("--target", required=True) - parser.add_argument("--op", required=True) - parser.add_argument("--operand-specs", required=True) - parser.add_argument("--context-attrs", default=None) - parser.add_argument( - "--method", - choices=("instantiate", "get_metadata"), - default="instantiate", - ) - parser.add_argument("--candidate-id", default=None) - args = parser.parse_args(argv) - - try: - operand_specs = json.loads(args.operand_specs) - context_attrs = json.loads(args.context_attrs) if args.context_attrs else {} - except json.JSONDecodeError as exc: - parser.error(f"invalid JSON input: {exc}") - - try: - client = DaemonClient(args.socket) - if args.method == "get_metadata": - result = client.get_metadata( - args.target, - args.op, - operand_specs, - context_attrs, - ) - sys.stdout.write(json.dumps(result)) - return - - result = client.instantiate( - args.target, - args.op, - operand_specs, - context_attrs, - args.candidate_id, - ) - except (DaemonError, OSError) as exc: - sys.stderr.write(f"Error: daemon RPC failed: {exc}\n") - raise SystemExit(1) from exc - - sys.stdout.write(result) - - -if __name__ == "__main__": - main() diff --git a/ptodsl/ptodsl/tilelib/serving/wire.py b/ptodsl/ptodsl/tilelib/serving/wire.py deleted file mode 100644 index 8f9185137b..0000000000 --- a/ptodsl/ptodsl/tilelib/serving/wire.py +++ /dev/null @@ -1,57 +0,0 @@ -# Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software, you can redistribute it and/or modify it under the terms and conditions of -# CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -# See LICENSE in the root of the software repository for the full text of the License. -"""Length-prefixed JSON framing for the TileLib daemon RPC.""" - -from __future__ import annotations - -import json - - -MAX_MESSAGE_SIZE = 64 * 1024 * 1024 - - -def recv_exactly(sock, length: int) -> bytes: - """Read exactly ``length`` bytes or fail if the peer closes early.""" - chunks = [] - remaining = length - while remaining: - chunk = sock.recv(remaining) - if not chunk: - raise ConnectionError("socket closed mid-message") - chunks.append(chunk) - remaining -= len(chunk) - return b"".join(chunks) - - -def send_message(sock, message: dict) -> None: - """Send one UTF-8 JSON message with a 4-byte big-endian length prefix.""" - payload = json.dumps(message).encode("utf-8") - if len(payload) > MAX_MESSAGE_SIZE: - raise ValueError( - f"message length {len(payload)} exceeds limit {MAX_MESSAGE_SIZE}" - ) - sock.sendall(len(payload).to_bytes(4, byteorder="big")) - sock.sendall(payload) - - -def recv_message(sock) -> dict: - """Receive one length-prefixed UTF-8 JSON message.""" - length = int.from_bytes(recv_exactly(sock, 4), byteorder="big") - if length > MAX_MESSAGE_SIZE: - raise ValueError( - f"message length {length} exceeds limit {MAX_MESSAGE_SIZE}" - ) - return json.loads(recv_exactly(sock, length).decode("utf-8")) - - -__all__ = [ - "MAX_MESSAGE_SIZE", - "recv_exactly", - "recv_message", - "send_message", -] diff --git a/ptodsl/tests/test_ptoas_cli.py b/ptodsl/tests/test_ptoas_cli.py index 51aa3c07dd..fd30b86f71 100644 --- a/ptodsl/tests/test_ptoas_cli.py +++ b/ptodsl/tests/test_ptoas_cli.py @@ -45,20 +45,13 @@ def test_launch_uses_standard_native_module_and_packaged_resources(self): self.assertEqual(exit_code, 0) native_module.main.assert_called_once_with( - [ - str(wrapper.resolve()), - "--ptodsl-pkg-path", - str(package_root.parent.resolve()), - "--tileops-pkg-path", - str(tileops_dir.parent.resolve()), - "--version", - ] + [str(wrapper.resolve()), "--version"] ) self.assertEqual(environment["PTOAS_BIN"], str(wrapper.resolve())) - self.assertEqual(environment["PTOAS_PYTHON_EXE"], _cli.sys.executable) + self.assertNotIn("PTOAS_PYTHON_EXE", environment) self.assertEqual(environment["PATH"], "/usr/bin") - def test_explicit_resource_options_are_not_overridden(self): + def test_user_arguments_are_forwarded_unchanged(self): with tempfile.TemporaryDirectory() as temp_dir: package_root = Path(temp_dir) / "install" / "ptoas" (package_root / "_runtime" / "share" / "ptoas" / "TileOps").mkdir( @@ -69,9 +62,7 @@ def test_explicit_resource_options_are_not_overridden(self): wrapper.write_text("", encoding="utf-8") native_module = self._make_native_module(package_root) arguments = [ - "--ptodsl-pkg-path=/custom/ptodsl", - "--tileops-pkg-path", - "/custom/tileops", + "--pto-arch=a5", "--version", ] @@ -92,11 +83,8 @@ def test_build_tree_uses_the_same_packaged_resource_layout(self): tileops_dir.mkdir(parents=True) native_module = self._make_native_module(package_root) - python_root, resolved_tileops = _cli._resolve_runtime_paths( - native_module - ) + resolved_tileops = _cli._resolve_tileops_dir(native_module) - self.assertEqual(python_root, package_root.parent.resolve()) self.assertEqual(resolved_tileops, tileops_dir.resolve()) def test_missing_tileops_resources_is_an_error(self): @@ -105,7 +93,7 @@ def test_missing_tileops_resources_is_an_error(self): native_module = self._make_native_module(package_root) with self.assertRaisesRegex(SystemExit, "TileOps"): - _cli._resolve_runtime_paths(native_module) + _cli._resolve_tileops_dir(native_module) if __name__ == "__main__": diff --git a/ptodsl/tests/test_tilelib_daemon.py b/ptodsl/tests/test_tilelib_daemon.py deleted file mode 100644 index 5d565d7a9c..0000000000 --- a/ptodsl/tests/test_tilelib_daemon.py +++ /dev/null @@ -1,266 +0,0 @@ -# Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software, you can redistribute it and/or modify it under the terms and conditions of -# CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -# See LICENSE in the root of the software repository for the full text of the License. -"""End-to-end tests for the PTODSL TileLib daemon's Unix-socket RPC.""" - -import os -import socket -import stat -import tempfile -import threading -import unittest - -from ptodsl.tilelib.serving.client import DaemonClient, DaemonError -from ptodsl.tilelib.serving.daemon import ( - TileLibDaemonServer, - _remove_socket_path, -) -from ptodsl.tilelib.serving.wire import MAX_MESSAGE_SIZE, recv_message - - -def _tile_spec(dtype="f32", shape=(8, 64)): - return { - "kind": "tile", - "dtype": dtype, - "shape": list(shape), - "valid_shape": list(shape), - "memory_space": "ub", - "config": { - "b_layout": "row_major", - "s_layout": "none_box", - "s_fractal_size": 512, - "pad_value": "0x0", - }, - } - - -def _view_spec(dtype="f32", shape=(1, 1, 1, 8, 64), strides=(512, 512, 512, 64, 1)): - return { - "kind": "view", - "dtype": dtype, - "shape": list(shape), - "strides": list(strides), - "memory_space": "gm", - } - - -# ExpandTileOp sends tadd as ins(src0, src1), outs(dst), matching the -# template parameter order (src0, src1, dst). -TADD_OPERANDS = [_tile_spec(), _tile_spec(), _tile_spec()] -TADD = "template_tadd" - - -class TileLibDaemonTest(unittest.TestCase): - def setUp(self): - self._temporary_directory = tempfile.TemporaryDirectory() - self.socket_path = os.path.join( - self._temporary_directory.name, - "ptodsl_tilelib.sock", - ) - self.server = TileLibDaemonServer(self.socket_path) - self._thread = threading.Thread( - target=self.server.serve_forever, - daemon=True, - ) - self._thread.start() - self.client = DaemonClient(self.socket_path) - - def tearDown(self): - self.server.shutdown() - self.server.server_close() - self._thread.join() - self._temporary_directory.cleanup() - - def test_ping(self): - self.assertEqual(self.client.ping(), "pong") - - def test_socket_is_accessible_only_by_owner(self): - mode = stat.S_IMODE(os.stat(self.socket_path).st_mode) - self.assertEqual(mode, 0o600) - - def test_instantiate_named_candidate_returns_structured_mlir(self): - mlir = self.client.instantiate( - "a5", - "pto.tadd", - TADD_OPERANDS, - candidate_id=TADD, - ) - self.assertIn(f"func.func @{TADD}", mlir) - for operation in ( - "pto.tile_buf_addr", - "!pto.ptr", - "pto.vlds", - "pto.vadd", - "pto.vsts", - "pto.plt_b32", - "pto.tilelang.instance", - ): - self.assertIn(operation, mlir) - self.assertNotIn("pto.castptr", mlir) - - def test_instantiate_uses_single_tadd_candidate_without_explicit_id(self): - mlir = self.client.instantiate("a5", "pto.tadd", TADD_OPERANDS) - self.assertIn(f"func.func @{TADD}", mlir) - - def test_get_metadata_returns_legal_candidates(self): - metadata = self.client.get_metadata("a5", "pto.tadd", TADD_OPERANDS) - candidates = metadata["candidates"] - self.assertEqual( - set(candidates), - {TADD}, - ) - - selected = candidates[TADD] - self.assertEqual(selected["loop_depth"], 2) - self.assertIsNone(selected["Tail"]) - self.assertFalse(selected["has_tail"]) - self.assertFalse(selected["is_post_update"]) - self.assertEqual(selected["iteration_axis"], "none") - self.assertEqual(selected["op_engine"], "vector") - self.assertEqual(selected["op_class"], "elementwise") - self.assertEqual(selected["tags"], ["elementwise", "binary"]) - - def test_cache_stats_and_clear_are_available_over_rpc(self): - arguments = ( - "a5", - "pto.tadd", - TADD_OPERANDS, - ) - self.client.instantiate( - *arguments, - candidate_id=TADD, - ) - self.client.instantiate( - *arguments, - candidate_id=TADD, - ) - - stats = self.client.get_stats() - self.assertEqual(stats["misses"], 1) - self.assertEqual(stats["hits"], 1) - self.assertEqual(stats["entries"], 1) - - self.assertEqual(self.client.clear(), {"cleared": True}) - self.assertEqual(self.client.get_stats()["entries"], 0) - - def test_cache_key_includes_context_attributes(self): - self.client.instantiate( - "a5", - "pto.tadd", - TADD_OPERANDS, - context_attrs={"variant": 0}, - candidate_id=TADD, - ) - self.client.instantiate( - "a5", - "pto.tadd", - TADD_OPERANDS, - context_attrs={"variant": 1}, - candidate_id=TADD, - ) - self.assertEqual(self.client.get_stats()["misses"], 2) - - def test_oversized_wire_message_is_rejected_before_payload_read(self): - receiver, sender = socket.socketpair() - self.addCleanup(receiver.close) - self.addCleanup(sender.close) - sender.sendall((MAX_MESSAGE_SIZE + 1).to_bytes(4, byteorder="big")) - - with self.assertRaisesRegex(ValueError, "exceeds limit"): - recv_message(receiver) - - def test_socket_cleanup_removes_broken_symlink(self): - missing_target = os.path.join( - self._temporary_directory.name, - "missing.sock", - ) - broken_link = os.path.join( - self._temporary_directory.name, - "broken.sock", - ) - os.symlink(missing_target, broken_link) - - _remove_socket_path(broken_link) - - self.assertFalse(os.path.lexists(broken_link)) - - def test_scalar_operand_template_instantiates(self): - operands = [ - _tile_spec(), - {"kind": "scalar", "dtype": "f32", "value": 1.0}, - _tile_spec(), - ] - - mlir = self.client.instantiate("a5", "pto.tadds", operands) - - self.assertIn("func.func @template_tadds", mlir) - self.assertIn("pto.vadds", mlir) - - def test_render_passes_context_attributes_into_template_body(self): - operands = [ - _tile_spec(dtype="f32", shape=(8, 64)), - _tile_spec(dtype="f32", shape=(8, 64)), - _tile_spec(dtype="i8", shape=(8, 64)), - ] - - mlir = self.client.instantiate( - "a5", - "pto.tcmp", - operands, - context_attrs={"cmp_mode": "gt"}, - candidate_id="template_tcmp", - ) - - self.assertIn('"gt"', mlir) - self.assertNotIn('"eq"', mlir) - - def test_vector_operand_metadata_is_accepted(self): - operands = [ - _tile_spec(), - _tile_spec(), - _tile_spec(), - _tile_spec(), - {"kind": "vector", "dtype": "i16", "shape": [4]}, - ] - - metadata = self.client.get_metadata("a5", "pto.tmrgsort", operands) - - self.assertIn("template_tmrgsort_multi_list2", metadata["candidates"]) - - def test_view_operand_template_instantiates(self): - operands = [_view_spec(), _tile_spec()] - - mlir = self.client.instantiate( - "a5", - "pto.tload", - operands, - candidate_id="template_tload_nd2nd", - ) - - self.assertIn("func.func @template_tload_nd2nd", mlir) - self.assertIn("pto.tensor_view_addr", mlir) - self.assertIn("pto.mte_gm_ub", mlir) - - def test_unsupported_operand_kind_is_rejected_explicitly(self): - operands = list(TADD_OPERANDS) - operands[0] = {"kind": "mystery", "dtype": "f32", "shape": [64]} - - with self.assertRaisesRegex(DaemonError, "supports tile, scalar, view, and vector operands"): - self.client.instantiate( - "a5", - "pto.tadd", - operands, - candidate_id=TADD, - ) - - def test_unknown_op_errors(self): - with self.assertRaises(DaemonError): - self.client.instantiate("a5", "pto.tnope", TADD_OPERANDS) - - -if __name__ == "__main__": - unittest.main() diff --git a/ptodsl/tests/test_tilelib_render.py b/ptodsl/tests/test_tilelib_render.py index ca5d4a677f..0f07d4668f 100644 --- a/ptodsl/tests/test_tilelib_render.py +++ b/ptodsl/tests/test_tilelib_render.py @@ -15,6 +15,8 @@ import unittest from pathlib import Path +from ptoas.mlir.dialects import pto as pto_dialect +from ptoas.mlir.ir import Context from ptodsl.tilelib import TileSpec, f32 from TileOps.a5.tadd import template_tadd @@ -65,6 +67,22 @@ def test_golden_fixture_uses_same_abstraction(self): for op in ("pto.tile_buf_addr", "!pto.ptr", "pto.vlds", "pto.vadd", "pto.vsts", "pto.plt_b32"): self.assertIn(op, golden) + def test_materialize_uses_borrowed_context_and_returns_fresh_modules(self): + context = Context() + pto_dialect.register_dialect(context, load=True) + spec = TileSpec(shape=(8, 64), dtype=f32) + artifact = template_tadd.specialize(src0=spec, src1=spec, dst=spec) + + first = artifact.materialize(context) + second = artifact.materialize(context) + + self.assertIs(first.context, context) + self.assertIs(second.context, context) + self.assertIsNot(first, second) + self.assertTrue(first.operation.verify()) + self.assertTrue(second.operation.verify()) + self.assertIn("func.func @template_tadd", artifact.mlir_text()) + if __name__ == "__main__": unittest.main() diff --git a/test/lit/vpto/expand_tile_op_ptodsl_tsub.pto b/test/lit/vpto/expand_tile_op_ptodsl_tsub.pto index 76819a40ee..b7c6b5374a 100644 --- a/test/lit/vpto/expand_tile_op_ptodsl_tsub.pto +++ b/test/lit/vpto/expand_tile_op_ptodsl_tsub.pto @@ -6,10 +6,10 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// Test that PTOAS can select the PTODSL TileLib daemon and expand the -// single-candidate pto.tsub template without using the legacy TileLang path. +// Test that PTOAS can use the in-process PTODSL TileLib service to expand the +// single-candidate pto.tsub template. // -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --tile-lib-backend=ptodsl %s -o - 2>/dev/null | FileCheck %s +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>/dev/null | FileCheck %s // CHECK: func.func @TSUB // CHECK-NOT: pto.tsub ins diff --git a/test/lit/vpto/expand_tile_op_ptodsl_view_stride_cache.pto b/test/lit/vpto/expand_tile_op_ptodsl_view_stride_cache.pto index 787b00ea8c..963d44f28d 100644 --- a/test/lit/vpto/expand_tile_op_ptodsl_view_stride_cache.pto +++ b/test/lit/vpto/expand_tile_op_ptodsl_view_stride_cache.pto @@ -10,7 +10,7 @@ // bodies. Two tstores with the same tile type but different destination view // strides must therefore not share one cached helper specialization. // -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --tile-lib-backend=ptodsl %s -o - 2>/dev/null | FileCheck %s +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>/dev/null | FileCheck %s // CHECK-LABEL: func.func @STORE_COMPACT // CHECK: %[[COMPACT_GM:.*]] = arith.constant 4 : i64 diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/run_ptoas_to_file.cmake b/test/tilelang_st/npu/a5/src/st/smoke/testcase/run_ptoas_to_file.cmake index 31b9b92438..b8a3a0070a 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/run_ptoas_to_file.cmake +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/run_ptoas_to_file.cmake @@ -28,17 +28,6 @@ endif() list(APPEND PTOAS_COMMAND --pto-backend=vpto) -set(PTOAS_TILE_LIB_BACKEND_EFFECTIVE "") -if(DEFINED PTOAS_TILE_LIB_BACKEND AND NOT PTOAS_TILE_LIB_BACKEND STREQUAL "") - set(PTOAS_TILE_LIB_BACKEND_EFFECTIVE "${PTOAS_TILE_LIB_BACKEND}") -elseif(DEFINED ENV{PTOAS_TILE_LIB_BACKEND} AND NOT "$ENV{PTOAS_TILE_LIB_BACKEND}" STREQUAL "") - set(PTOAS_TILE_LIB_BACKEND_EFFECTIVE "$ENV{PTOAS_TILE_LIB_BACKEND}") -endif() - -if(NOT PTOAS_TILE_LIB_BACKEND_EFFECTIVE STREQUAL "") - list(APPEND PTOAS_COMMAND "--tile-lib-backend=${PTOAS_TILE_LIB_BACKEND_EFFECTIVE}") -endif() - if(PTOAS_ENABLE_INSERT_SYNC) list(APPEND PTOAS_COMMAND --enable-insert-sync) endif() diff --git a/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake b/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake index 31b9b92438..b8a3a0070a 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake +++ b/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake @@ -28,17 +28,6 @@ endif() list(APPEND PTOAS_COMMAND --pto-backend=vpto) -set(PTOAS_TILE_LIB_BACKEND_EFFECTIVE "") -if(DEFINED PTOAS_TILE_LIB_BACKEND AND NOT PTOAS_TILE_LIB_BACKEND STREQUAL "") - set(PTOAS_TILE_LIB_BACKEND_EFFECTIVE "${PTOAS_TILE_LIB_BACKEND}") -elseif(DEFINED ENV{PTOAS_TILE_LIB_BACKEND} AND NOT "$ENV{PTOAS_TILE_LIB_BACKEND}" STREQUAL "") - set(PTOAS_TILE_LIB_BACKEND_EFFECTIVE "$ENV{PTOAS_TILE_LIB_BACKEND}") -endif() - -if(NOT PTOAS_TILE_LIB_BACKEND_EFFECTIVE STREQUAL "") - list(APPEND PTOAS_COMMAND "--tile-lib-backend=${PTOAS_TILE_LIB_BACKEND_EFFECTIVE}") -endif() - if(PTOAS_ENABLE_INSERT_SYNC) list(APPEND PTOAS_COMMAND --enable-insert-sync) endif() diff --git a/test/tilelang_st/script/run_a5_st_all_parallel.py b/test/tilelang_st/script/run_a5_st_all_parallel.py index 8299c743f5..962636df1e 100755 --- a/test/tilelang_st/script/run_a5_st_all_parallel.py +++ b/test/tilelang_st/script/run_a5_st_all_parallel.py @@ -118,15 +118,11 @@ def _run_one(job, args, ptoas_bin, output_root, base_env): build_dir = job_root / "build" tmp_dir = job_root / "tmp" log_path = output_root / "logs" / f"{job_name}.log" - socket_path = Path("/tmp") / f"ptoas_st_{kind}_{testcase}_{os.getpid()}.sock" started = time.time() env = base_env.copy() env["TMPDIR"] = str(tmp_dir) env["PTODSL_CACHE_DIR"] = str(job_root / "ptodsl-cache") - env["PTOAS_DAEMON_SOCKET_PATH"] = str(socket_path) - if args.tile_lib_backend: - env["PTOAS_TILE_LIB_BACKEND"] = args.tile_lib_backend tmp_dir.mkdir(parents=True, exist_ok=True) (output_root / "logs").mkdir(parents=True, exist_ok=True) @@ -139,7 +135,6 @@ def _run_one(job, args, ptoas_bin, output_root, base_env): "seconds": 0.0, "log": str(log_path), "build_dir": str(build_dir), - "socket": str(socket_path), } try: @@ -148,10 +143,7 @@ def _run_one(job, args, ptoas_bin, output_root, base_env): log_handle.write(f"# testcase: {testcase}\n") log_handle.write(f"# source: {job['target_dir']}\n") log_handle.write(f"# build: {build_dir}\n") - log_handle.write(f"# PTOAS_DAEMON_SOCKET_PATH={socket_path}\n") log_handle.write(f"# PTODSL_CACHE_DIR={env['PTODSL_CACHE_DIR']}\n") - if args.tile_lib_backend: - log_handle.write(f"# PTOAS_TILE_LIB_BACKEND={args.tile_lib_backend}\n") log_handle.write("\n") cmake_cmd = [ @@ -164,10 +156,7 @@ def _run_one(job, args, ptoas_bin, output_root, base_env): f"-DSOC_VERSION={DEFAULT_SOC_VERSION}", f"-DTEST_CASE={testcase}", f"-DPTOAS_BIN={ptoas_bin}", - f"-DPTOAS_DAEMON_SOCKET_PATH={socket_path}", ] - if args.tile_lib_backend: - cmake_cmd.append(f"-DPTOAS_TILE_LIB_BACKEND={args.tile_lib_backend}") rc = _run_logged(cmake_cmd, log_handle, output_root, env) if rc == 0: @@ -270,11 +259,6 @@ def _parse_args(): default=str(_default_output_root(repo_root)), help="Directory for logs, summaries, and per-testcase build trees.", ) - parser.add_argument( - "--tile-lib-backend", - default=os.environ.get("PTOAS_TILE_LIB_BACKEND", ""), - help="Optional PTOAS tile-lib backend, for example ptodsl.", - ) parser.add_argument("--full-only", action="store_true", help="Run only non-smoke ST cases.") parser.add_argument("--smoke-only", action="store_true", help="Run only smoke ST cases.") parser.add_argument("--list", action="store_true", help="List selected jobs and exit.") @@ -339,9 +323,7 @@ def main(): print(f"[INFO] run_mode={args.run_mode} soc={SOC_VERSION} ({DEFAULT_SOC_VERSION})") print(f"[INFO] ptoas={ptoas_bin}") print(f"[INFO] output_root={output_root}") - if args.tile_lib_backend: - print(f"[INFO] PTOAS_TILE_LIB_BACKEND={args.tile_lib_backend}") - print("[INFO] each testcase uses its own build dir, PTODSL cache, TMPDIR, and daemon socket") + print("[INFO] each testcase uses its own build dir, PTODSL cache, and TMPDIR") results = [] max_workers = min(args.jobs, len(jobs)) diff --git a/test/tilelang_st/script/run_ptodsl_st_parallel.py b/test/tilelang_st/script/run_ptodsl_st_parallel.py index 4c1e4a0ac7..045cb94ab8 100755 --- a/test/tilelang_st/script/run_ptodsl_st_parallel.py +++ b/test/tilelang_st/script/run_ptodsl_st_parallel.py @@ -7,7 +7,7 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""Run TileLang ST testcases with the PTODSL TileLib backend and per-test logs.""" +"""Run TileLang ST testcases with PTODSL TileLib and per-test logs.""" import argparse import concurrent.futures @@ -69,7 +69,6 @@ def _run_build(args, default_soc_version, target_dir, log_dir, ptoas_bin): started = time.time() with build_log.open("w", encoding="utf-8") as handle: handle.write(f"# cwd: {target_dir}\n") - handle.write(f"# PTOAS_TILE_LIB_BACKEND={args.tile_lib_backend}\n") handle.write(f"# ptoas: {ptoas_bin}\n") handle.write("# command: run_st.build_project(..., testcase='all', ...)\n\n") handle.flush() @@ -116,13 +115,11 @@ def _run_one_testcase(args, testcase, target_dir, log_dir, ptoas_bin): command.append("--smoke") env = os.environ.copy() - env["PTOAS_TILE_LIB_BACKEND"] = args.tile_lib_backend started = time.time() with log_path.open("w", encoding="utf-8") as handle: handle.write(f"# testcase: {testcase}\n") handle.write(f"# cwd: {target_dir}\n") - handle.write(f"# PTOAS_TILE_LIB_BACKEND={args.tile_lib_backend}\n") handle.write("# command: " + " ".join(command) + "\n\n") handle.flush() proc = subprocess.Popen( @@ -150,10 +147,7 @@ def _run_one_testcase(args, testcase, target_dir, log_dir, ptoas_bin): def _parse_args(): repo_root = _repo_root() parser = argparse.ArgumentParser( - description=( - "Run TileLang ST testcases in parallel with PTOAS_TILE_LIB_BACKEND=ptodsl " - "and save one log per testcase." - ) + description="Run TileLang ST testcases in parallel and save one log per testcase." ) parser.add_argument("-r", "--run-mode", default="sim", help="Run mode: sim or npu.") parser.add_argument("-v", "--soc-version", default="a5", help="SoC version key, default: a5.") @@ -182,11 +176,6 @@ def _parse_args(): default=None, help="Directory for build.log, one .log per testcase, and summary files.", ) - parser.add_argument( - "--tile-lib-backend", - default="ptodsl", - help="Value for PTOAS_TILE_LIB_BACKEND, default: ptodsl.", - ) parser.add_argument( "--full", action="store_true", @@ -250,14 +239,12 @@ def main(): print(f"[INFO] target_dir={target_dir}") print(f"[INFO] ptoas={ptoas_bin}") print(f"[INFO] logs={log_dir}") - print(f"[INFO] PTOAS_TILE_LIB_BACKEND={args.tile_lib_backend}") if args.dry_run: for testcase in selected: print(f"[DRY-RUN] {testcase}") return 0 - os.environ["PTOAS_TILE_LIB_BACKEND"] = args.tile_lib_backend default_soc_version = run_all_st.SOC_VERSION_MAP[args.soc_version] results = [] @@ -311,7 +298,6 @@ def main(): break summary = { - "backend": args.tile_lib_backend, "run_mode": args.run_mode, "soc_version": args.soc_version, "smoke": args.smoke, diff --git a/tools/ptoas/CMakeLists.txt b/tools/ptoas/CMakeLists.txt index 2ce4929432..d2aadd3c00 100644 --- a/tools/ptoas/CMakeLists.txt +++ b/tools/ptoas/CMakeLists.txt @@ -16,7 +16,6 @@ set(PTOAS_RUNTIME_SOURCES driver.cpp VPTOHostStubEmission.cpp ObjectEmission.cpp - TilelangDaemon.cpp ) add_library(PTOASVFSIMTSizePatcher STATIC @@ -67,13 +66,6 @@ endforeach() function(ptoas_configure_runtime_compile_target target_name) target_compile_definitions(${target_name} PRIVATE PTOAS_RELEASE_VERSION="${PTOAS_CLI_VERSION}" - # Source-tree defaults for TileLib expansion. These let ptoas run directly - # from the build tree without passing --ptodsl-pkg-path / - # --tileops-pkg-path. Installed layouts are expected - # to launch through the wrapper/launcher flow, which injects explicit - # runtime paths instead of relying on executable-path probing. - PTOAS_DEFAULT_PTODSL_PKG_PATH="${CMAKE_SOURCE_DIR}/ptodsl" - PTOAS_DEFAULT_TILEOPS_PKG_PATH="${CMAKE_SOURCE_DIR}/lib" ${ARGN} ) add_dependencies(${target_name} diff --git a/tools/ptoas/NativeModule.cpp b/tools/ptoas/NativeModule.cpp index 5a474cbd27..36b12ce341 100644 --- a/tools/ptoas/NativeModule.cpp +++ b/tools/ptoas/NativeModule.cpp @@ -9,7 +9,12 @@ #include "ptoas.h" #include "PTOModule.h" +#include "PTO/Transforms/TileLibService.h" +#include "mlir/Bindings/Python/PybindAdaptors.h" +#include "mlir/CAPI/IR.h" + +#include "llvm/Support/raw_ostream.h" #include "pybind11/pybind11.h" #include "pybind11/stl.h" @@ -20,6 +25,87 @@ namespace py = pybind11; namespace { +class PythonTileLibService final : public mlir::pto::TileLibService { +public: + explicit PythonTileLibService(py::object contextOwner) + : contextOwner(std::move(contextOwner)) {} + + mlir::FailureOr + getMetadata(const mlir::pto::TileLibMaterializationRequest &request) override { + py::gil_scoped_acquire acquire; + try { + return py::cast(getRuntime().attr("metadata")( + request.target, request.op, request.operandSpecsJson, + request.contextAttrsJson)); + } catch (const py::error_already_set &error) { + llvm::errs() << "TileLib: PTODSL metadata query raised Python " + "exception:\n" + << error.what() << "\n"; + return mlir::failure(); + } + } + + mlir::FailureOr + materialize(const mlir::pto::TileLibMaterializationRequest &request, + mlir::MLIRContext &context) override { + py::gil_scoped_acquire acquire; + try { + MlirContext pythonContext = py::cast(contextOwner); + if (unwrap(pythonContext) != &context) { + llvm::errs() << "TileLib: Python context does not match the PTOAS " + "MLIRContext\n"; + return mlir::failure(); + } + + py::tuple result = getRuntime().attr("materialize")( + request.target, request.op, request.operandSpecsJson, + request.contextAttrsJson, request.candidateId, contextOwner); + if (result.size() != 2) + throw py::value_error( + "PTODSL materialize() must return (module, entry_symbol)"); + + // MlirModule is a non-owning handle. Keep result[0] alive until the + // complete source module has been cloned into C++ ownership. + py::object moduleOwner = result[0]; + MlirModule rawModule = py::cast(moduleOwner); + if (!mlirContextEqual(mlirModuleGetContext(rawModule), pythonContext)) { + llvm::errs() << "TileLib: PTODSL returned a module from a different " + "MLIRContext\n"; + return mlir::failure(); + } + + mlir::ModuleOp source = unwrap(rawModule); + auto cloned = mlir::cast(source->clone()); + mlir::pto::TileLibMaterialization materialization{ + mlir::OwningOpRef(cloned), + py::cast(result[1])}; + return materialization; + } catch (const py::error_already_set &error) { + llvm::errs() << "TileLib: PTODSL materialization raised Python " + "exception:\n" + << error.what() << "\n"; + return mlir::failure(); + } catch (const std::exception &error) { + llvm::errs() << "TileLib: invalid PTODSL materialization result: " + << error.what() << "\n"; + return mlir::failure(); + } + } + +private: + py::object &getRuntime() { + if (!runtime) + runtime = py::module_::import("ptodsl.tilelib._compiler_runtime"); + return runtime; + } + + // These objects are created and destroyed by runPTOASFromPython while the + // calling thread owns the GIL. materialize() reacquires it for every DSL + // invocation because the native compiler releases it around the driver. + py::object contextOwner; + py::object runtime; +}; + int runPTOASFromPython(const std::vector &arguments) { std::vector storage = arguments; std::vector argv; @@ -27,8 +113,19 @@ int runPTOASFromPython(const std::vector &arguments) { for (std::string &argument : storage) argv.push_back(argument.data()); - py::gil_scoped_release release; - return mlir::pto::runPTOAS(static_cast(argv.size()), argv.data()); + py::object contextOwner = + py::module_::import("ptoas.mlir.ir").attr("Context")(); + MlirContext rawContext = py::cast(contextOwner); + auto tileLibService = + std::make_shared(contextOwner); + + int result; + { + py::gil_scoped_release release; + result = mlir::pto::runPTOAS(static_cast(argv.size()), argv.data(), + *unwrap(rawContext), tileLibService); + } + return result; } } // namespace diff --git a/tools/ptoas/TilelangDaemon.cpp b/tools/ptoas/TilelangDaemon.cpp deleted file mode 100644 index 9ccc86221c..0000000000 --- a/tools/ptoas/TilelangDaemon.cpp +++ /dev/null @@ -1,154 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -#include "PTO/Support/PythonExecutable.h" -#include "TilelangDaemon.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/Support/FileSystem.h" -#include "llvm/Support/Program.h" -#include -#include -#include -#include -#include -#include - -extern char **environ; - -namespace ptoas { - -std::optional> DaemonManager::processInfo; - -std::string DaemonManager::generateSocketPath() { - return "/tmp/tilelib_daemon_" + std::to_string(::getpid()) + ".sock"; -} - -bool DaemonManager::start(const std::string &socketPath, - const std::string &daemonModule, - const std::string &pythonExe, - const std::string &pkgPath, - const std::string &templateDir) { - auto pythonPath = - mlir::pto::resolvePythonExecutable(pythonExe.empty() ? "python3" - : pythonExe); - if (!pythonPath) { - llvm::errs() << "Error: Cannot find Python executable '" - << (pythonExe.empty() ? "python3" : pythonExe) - << "' for daemon\n"; - return false; - } - - llvm::SmallVector args = { - *pythonPath, "-m", daemonModule, "--socket", socketPath, - }; - if (!templateDir.empty()) { - args.push_back("--template-dir"); - args.push_back(templateDir); - } - - llvm::SmallVector envp; - std::string pythonPathEnv; - std::vector envStorage; - - if (!pkgPath.empty()) { - const char *existingPath = ::getenv("PYTHONPATH"); - pythonPathEnv = "PYTHONPATH=" + pkgPath; - if (existingPath && existingPath[0] != '\0') { - pythonPathEnv += ":"; - pythonPathEnv += existingPath; - } - for (char **e = environ; *e; ++e) { - llvm::StringRef entry(*e); - if (entry.starts_with("PYTHONPATH=")) - continue; - envStorage.push_back(std::string(entry)); - } - envStorage.push_back(pythonPathEnv); - for (auto &s : envStorage) - envp.push_back(s); - } - - std::string errMsg; - bool executionFailed = false; - - llvm::sys::ProcessInfo procInfo = llvm::sys::ExecuteNoWait( - *pythonPath, args, - !pkgPath.empty() - ? std::optional>(envp) - : std::nullopt, - {}, 0, &errMsg, &executionFailed, nullptr, true); - - if (executionFailed || procInfo.Pid == llvm::sys::ProcessInfo::InvalidPid) { - llvm::errs() << "Error: Failed to start TileLib daemon module '" - << daemonModule << "': " << errMsg << "\n"; - return false; - } - - processInfo = std::make_pair(procInfo.Pid, socketPath); - - // Python startup time depends on the selected TileLib frontend and its - // imports. Poll instead of relying on one fixed sleep. - bool socketReady = false; - // PTODSL imports can be noticeably slower on heavily loaded CI runners where - // many ptoas processes start TileLib daemons concurrently. - for (int attempt = 0; attempt < 600; ++attempt) { - std::this_thread::sleep_for(std::chrono::milliseconds(50)); - if (llvm::sys::fs::exists(socketPath)) { - socketReady = true; - break; - } - } - - if (!socketReady) { - llvm::errs() << "Error: Daemon socket not created at " << socketPath << "\n"; - llvm::errs() << "Note: Daemon process started (pid=" << procInfo.Pid - << ") but socket not found. Check daemon logs.\n"; - kill(procInfo.Pid, SIGTERM); - processInfo = std::nullopt; - return false; - } - - llvm::errs() << "TileLib daemon '" << daemonModule << "' started (pid=" - << procInfo.Pid - << ", socket=" << socketPath << ")\n"; - return true; -} - -void DaemonManager::stop() { - if (!processInfo) - return; - - int pid = processInfo->first; - std::string socketPath = processInfo->second; - - kill(pid, SIGTERM); - - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - - if (llvm::sys::fs::exists(socketPath)) { - llvm::sys::fs::remove(socketPath); - } - - llvm::errs() << "TileLib daemon stopped (pid=" << pid << ")\n"; - processInfo = std::nullopt; -} - -bool DaemonManager::isRunning() { - return processInfo.has_value(); -} - -static void daemonCleanupHandler() { - DaemonManager::stop(); -} - -void registerDaemonCleanup() { - std::atexit(daemonCleanupHandler); -} - -} // namespace ptoas diff --git a/tools/ptoas/TilelangDaemon.h b/tools/ptoas/TilelangDaemon.h deleted file mode 100644 index 8b369f6fba..0000000000 --- a/tools/ptoas/TilelangDaemon.h +++ /dev/null @@ -1,44 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -#ifndef PTOAS_TILELANG_DAEMON_H -#define PTOAS_TILELANG_DAEMON_H - -#include -#include -#include - -namespace llvm::sys { -using procid_t = int; -} - -namespace ptoas { - -class DaemonManager { -public: - static std::string generateSocketPath(); - - static bool start(const std::string &socketPath, - const std::string &daemonModule, - const std::string &pythonExe, - const std::string &pkgPath, - const std::string &templateDir = ""); - - static void stop(); - - static bool isRunning(); - -private: - static std::optional> processInfo; -}; - -void registerDaemonCleanup(); - -} // namespace ptoas - -#endif // PTOAS_TILELANG_DAEMON_H diff --git a/tools/ptoas/driver.cpp b/tools/ptoas/driver.cpp index 4c44d980ce..5f4320b992 100644 --- a/tools/ptoas/driver.cpp +++ b/tools/ptoas/driver.cpp @@ -678,8 +678,16 @@ static LogicalResult emitVPTOLLVMFatobj( mlir::pto::PTOASContext::PTOASContext(DialectRegistry ®istry, llvm::StringRef outputPath, int argc, char **argv) - : mlirContext(registry), outputPath(outputPath.str()), argc(argc), - argv(argv) {} + : ownedMlirContext(std::make_unique(registry)), + mlirContext(ownedMlirContext.get()), outputPath(outputPath.str()), + argc(argc), argv(argv) {} + +mlir::pto::PTOASContext::PTOASContext( + MLIRContext &borrowedContext, + std::shared_ptr tileLibService, + llvm::StringRef outputPath, int argc, char **argv) + : mlirContext(&borrowedContext), tileLibService(std::move(tileLibService)), + outputPath(outputPath.str()), argc(argc), argv(argv) {} mlir::pto::PTOASContext::~PTOASContext() = default; @@ -694,11 +702,16 @@ mlir::pto::PTOASContext::initializeEnvironment(bool requiresToolchain, void mlir::pto::PTOASContext::initializeMLIRContext() { // Be tolerant: ptobc decode may materialize ops from dialects that aren't // explicitly registered/loaded in this tool yet. - mlirContext.allowUnregisteredDialects(true); - mlir::pto::loadPTOASDialects(mlirContext); + mlirContext->allowUnregisteredDialects(true); + mlir::pto::loadPTOASDialects(*mlirContext); } -MLIRContext &mlir::pto::PTOASContext::getMLIRContext() { return mlirContext; } +MLIRContext &mlir::pto::PTOASContext::getMLIRContext() { return *mlirContext; } + +std::shared_ptr +mlir::pto::PTOASContext::getTileLibService() const { + return tileLibService; +} void mlir::pto::PTOASContext::setArch(std::string value) { arch = std::move(value); @@ -1266,9 +1279,13 @@ static LogicalResult writeTextOutput(llvm::StringRef output, // +-------------+ +------------------------------------------+ // | C++ source | | fatobj | // +-------------+ +------------------------------------------+ -static int runPTOASDriver(int argc, char **argv) { +static int runPTOASDriver( + int argc, char **argv, MLIRContext *borrowedContext = nullptr, + std::shared_ptr tileLibService = nullptr) { DialectRegistry registry; mlir::pto::registerPTOASDialects(registry); + if (borrowedContext) + borrowedContext->appendDialectRegistry(registry); mlir::pto::registerPTOASPassesAndCLOptions(); llvm::cl::SetVersionPrinter(printPTOASVersion); @@ -1283,10 +1300,17 @@ static int runPTOASDriver(int argc, char **argv) { llvm::errs())) return 1; - PTOASContext context(registry, outputFilename, argc, argv); - context.setOutputCANNVersionOverride(outputCANNVersionOverride); - context.setVFSIMTSizeFixMode(mlir::pto::vptoFixVFSIMTSize); - context.initializeMLIRContext(); + std::unique_ptr context; + if (borrowedContext) { + context = std::make_unique( + *borrowedContext, std::move(tileLibService), outputFilename, argc, argv); + } else { + context = + std::make_unique(registry, outputFilename, argc, argv); + } + context->setOutputCANNVersionOverride(outputCANNVersionOverride); + context->setVFSIMTSizeFixMode(mlir::pto::vptoFixVFSIMTSize); + context->initializeMLIRContext(); std::unique_ptr inputBuffer = readInputBuffer(); if (!inputBuffer) @@ -1294,24 +1318,24 @@ static int runPTOASDriver(int argc, char **argv) { std::string arch; OwningOpRef module = loadInputModule( - std::move(inputBuffer), context.getMLIRContext(), cliArchSpecified, arch); + std::move(inputBuffer), context->getMLIRContext(), cliArchSpecified, arch); if (!module) return 1; - context.setArch(std::move(arch)); + context->setArch(std::move(arch)); mlir::pto::BackendInfo backendInfo; if (failed(buildBackendInfo(module.get(), cliBackendSpecified, backendInfo))) return 1; - context.setBackendInfo(std::move(backendInfo)); - (void)context.initializeEnvironment(context.getBackendInfo().requiresToolchain, - llvm::errs()); + context->setBackendInfo(std::move(backendInfo)); + (void)context->initializeEnvironment( + context->getBackendInfo().requiresToolchain, llvm::errs()); mlir::pto::PTOASCompileResult result; - if (failed(runPTOASJobs(module, context, result))) + if (failed(runPTOASJobs(module, *context, result))) return 1; if (result.kind == mlir::pto::PTOASCompileResultKind::Text) - return failed(writeTextOutput(result.textOutput, context.getOutputPath())); + return failed(writeTextOutput(result.textOutput, context->getOutputPath())); if (result.kind == mlir::pto::PTOASCompileResultKind::MixedObject) return 0; @@ -1322,3 +1346,10 @@ static int runPTOASDriver(int argc, char **argv) { int mlir::pto::runPTOAS(int argc, char **argv) { return runPTOASDriver(argc, argv); } + +int mlir::pto::runPTOAS( + int argc, char **argv, MLIRContext &borrowedContext, + std::shared_ptr tileLibService) { + return runPTOASDriver(argc, argv, &borrowedContext, + std::move(tileLibService)); +} diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 8fdaa31afa..089e079108 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -14,7 +14,6 @@ #include "PTO/Transforms/Passes.h" #include "PTO/Transforms/BufferizableOpInterfaceImpl.h" #include "VPTOHostStubEmission.h" -#include "TilelangDaemon.h" #include "PTO/Transforms/CppPostprocess.h" #include "mlir/AsmParser/AsmParserState.h" #include "mlir/IR/MLIRContext.h" @@ -230,48 +229,6 @@ void mlir::pto::loadPTOASDialects(MLIRContext &context) { context.getOrLoadDialect(); } -static bool pathExists(llvm::StringRef path) { - return !path.empty() && llvm::sys::fs::exists(path); -} - -static std::string getParentDir(llvm::StringRef path) { - llvm::SmallString<256> parent(path); - llvm::sys::path::remove_filename(parent); - llvm::sys::path::remove_dots(parent, true); - return std::string(parent); -} - -static std::string joinPath(llvm::StringRef lhs, llvm::StringRef rhs) { - llvm::SmallString<256> joined(lhs); - llvm::sys::path::append(joined, rhs); - llvm::sys::path::remove_dots(joined, true); - return std::string(joined); -} - -static std::string detectInstalledPythonPkgRoot(const char *argv0, - llvm::StringRef packageName) { - std::string exePath = llvm::sys::fs::getMainExecutable(argv0, (void *)&main); - if (exePath.empty()) - return {}; - - const std::string exeDir = getParentDir(exePath); - const std::string prefixDir = getParentDir(exeDir); - const std::string installedPkg = joinPath(prefixDir, packageName); - if (pathExists(installedPkg)) - return prefixDir; - return {}; -} - -static bool hasCLIOption(int argc, char **argv, llvm::StringRef option) { - const std::string optionWithValue = (option + "=").str(); - for (int i = 1; i < argc; ++i) { - llvm::StringRef arg(argv[i]); - if (arg == option || arg.starts_with(optionWithValue)) - return true; - } - return false; -} - static LogicalResult applyConfiguredPassManagerCLOptions( PassManager &pm, llvm::StringRef pipelineName, llvm::raw_ostream &diagOS = llvm::errs()) { @@ -453,128 +410,6 @@ static llvm::cl::opt enableTileOpExpand( "--pto-backend=vpto."), llvm::cl::init(false)); -#ifndef PTOAS_DEFAULT_PTODSL_PKG_PATH -#define PTOAS_DEFAULT_PTODSL_PKG_PATH "" -#endif -#ifndef PTOAS_DEFAULT_TILEOPS_PKG_PATH -#define PTOAS_DEFAULT_TILEOPS_PKG_PATH "" -#endif - -static llvm::cl::opt ptodslPkgPath( - "ptodsl-pkg-path", - llvm::cl::desc("PYTHONPATH for the ptodsl package " - "(default: /ptodsl, baked in at build time)"), - llvm::cl::init(PTOAS_DEFAULT_PTODSL_PKG_PATH)); - -static llvm::cl::opt tileopsPkgPath( - "tileops-pkg-path", - llvm::cl::desc("PYTHONPATH for the TileOps PTODSL template package " - "(default: /lib, baked in at build time)"), - llvm::cl::init(PTOAS_DEFAULT_TILEOPS_PKG_PATH)); - -static llvm::cl::opt daemonSocketPath( - "daemon-socket-path", - llvm::cl::desc("Path to Unix domain socket for daemon RPC " - "(default: /tmp/tilelib_daemon_{pid}.sock)"), - llvm::cl::init("")); - -enum class TileLibBackend { - PTODSL, -}; - -static llvm::cl::opt tileLibBackend( - "tile-lib-backend", - llvm::cl::desc("TileLib backend used by ExpandTileOp"), - llvm::cl::values( - clEnumValN(TileLibBackend::PTODSL, "ptodsl", - "Use the PTODSL TileLib daemon")), - llvm::cl::init(TileLibBackend::PTODSL)); - -static std::string resolveTileLibPythonExe() { - const char *pythonExe = ::getenv("PTOAS_PYTHON_EXE"); - if (pythonExe && pythonExe[0] != '\0') - return pythonExe; - return "python3"; -} - -static pto::ExpandTileOpOptions resolveExpandTileOpOptions(int argc, - char **argv) { - pto::ExpandTileOpOptions expandOpts; - expandOpts.pythonExe = resolveTileLibPythonExe(); - std::string resolvedPtodslPkgPath = ptodslPkgPath; - std::string resolvedTileOpsPkgPath = tileopsPkgPath; - - if (!hasCLIOption(argc, argv, "--ptodsl-pkg-path")) { - const char *envPtodslRoot = ::getenv("PTODSL_PYTHON_ROOT"); - if (envPtodslRoot && envPtodslRoot[0] != '\0') - resolvedPtodslPkgPath = envPtodslRoot; - else { - std::string installedPtodslPkgPath = - detectInstalledPythonPkgRoot(argv[0], "ptodsl"); - if (!installedPtodslPkgPath.empty()) - resolvedPtodslPkgPath = installedPtodslPkgPath; - } - } - - if (!hasCLIOption(argc, argv, "--tileops-pkg-path")) { - const char *envTileOpsRoot = ::getenv("PTO_TILEOPS_PYTHON_ROOT"); - if (envTileOpsRoot && envTileOpsRoot[0] != '\0') - resolvedTileOpsPkgPath = envTileOpsRoot; - else { - std::string installedTileOpsPkgPath = - detectInstalledPythonPkgRoot(argv[0], "TileOps"); - if (!installedTileOpsPkgPath.empty()) - resolvedTileOpsPkgPath = installedTileOpsPkgPath; - } - } - - expandOpts.tileLibBackend = "ptodsl"; - expandOpts.daemonHelperModule = "ptodsl.tilelib.serving.helper"; - expandOpts.tileLibPkgPath = resolvedPtodslPkgPath; - if (!resolvedTileOpsPkgPath.empty()) { - if (!expandOpts.tileLibPkgPath.empty()) - expandOpts.tileLibPkgPath += ":"; - expandOpts.tileLibPkgPath += resolvedTileOpsPkgPath; - } - - // Daemon mode is default (no CLI option needed) - // Automatically start daemon for instance caching - std::string socket = daemonSocketPath; - if (socket.empty()) - socket = ptoas::DaemonManager::generateSocketPath(); - - // Register cleanup handler (daemon will be stopped on PTOAS exit) - ptoas::registerDaemonCleanup(); - - // Try to start daemon automatically - if (ptoas::DaemonManager::start(socket, "ptodsl.tilelib.serving.daemon", - expandOpts.pythonExe, - expandOpts.tileLibPkgPath, "")) { - expandOpts.daemonSocketPath = socket; - llvm::errs() << "Info: " << expandOpts.tileLibBackend - << " TileLib daemon started successfully\n"; - } else { - expandOpts.daemonSocketPath = ""; - llvm::errs() - << "Error: Failed to start the PTODSL TileLib daemon; no TileLang " - "fallback will be used\n"; - } - - return expandOpts; -} - - -static pto::InsertTemplateAttributesOptions -buildInsertTemplateAttributesOptions( - const pto::ExpandTileOpOptions &expandOptions) { - pto::InsertTemplateAttributesOptions options; - options.pythonExe = expandOptions.pythonExe; - options.daemonSocketPath = expandOptions.daemonSocketPath; - options.tileLibPkgPath = expandOptions.tileLibPkgPath; - options.daemonHelperModule = expandOptions.daemonHelperModule; - return options; -} - static llvm::cl::opt enableOpFusion( "enable-op-fusion", llvm::cl::desc("Control A5 tile fusion on level2/level3. Disabled by " @@ -2955,7 +2790,7 @@ static void prepareVPTOForEmission(PassManager &pm) { static void lowerPTOToVPTOBackend(PassManager &pm, ModuleOp module, - const pto::ExpandTileOpOptions &expandOpts) { + std::shared_ptr tileLibService) { auto &kernelModulePM = pm.nest(); auto moduleArchAttr = module->getAttrOfType("pto.target_arch"); @@ -2972,7 +2807,8 @@ lowerPTOToVPTOBackend(PassManager &pm, ModuleOp module, return; } - kernelModulePM.addPass(pto::createExpandTileOpPass(expandOpts)); + kernelModulePM.addPass( + pto::createExpandTileOpPass(std::move(tileLibService))); kernelModulePM.addPass(pto::createPTOInlineLibCallPass()); kernelModulePM.addNestedPass( @@ -3068,20 +2904,14 @@ static int emitVPTOBackendResult(ModuleOp module, PTOASCompileResult &result, static LogicalResult runVPTOBackendPipeline(OwningOpRef &module, bool hasTileOpsToExpand, - const pto::ExpandTileOpOptions - *expandOptions) { + std::shared_ptr + tileLibService) { PassManager pm(module->getContext()); pm.enableVerifier(); pm.addPass(pto::createVPTOSplitCVModulePass()); pm.addPass(pto::createVPTONormalizeContainerPass()); - if (hasTileOpsToExpand) { - if (!expandOptions) { - llvm::errs() << "Error: tile expansion requires resolved TileLib " - "options.\n"; - return failure(); - } - lowerPTOToVPTOBackend(pm, module.get(), *expandOptions); - } + if (hasTileOpsToExpand) + lowerPTOToVPTOBackend(pm, module.get(), std::move(tileLibService)); auto &kernelModulePM = pm.nest(); // Inline legal direct calls before VMI layout assignment so private helper // bodies participate in one caller-local layout decision. The Func @@ -3150,8 +2980,6 @@ int mlir::pto::compilePTOASModule( return 1; std::string arch = resolveEffectiveTargetArch(*module, context.getArch()); - int argc = context.getArgc(); - char **argv = context.getArgv(); // Name-hint provenance: textual .pto inputs had their SSA/arg/block-arg names // attached to op Locations by the driver right after parsing. Collect the @@ -3377,10 +3205,6 @@ int mlir::pto::compilePTOASModule( } const bool hasTileOpsToExpand = hasUnexpandedTileOps(*module); - std::optional expandOptions; - if (effectiveBackend == PTOBackend::VPTO && hasTileOpsToExpand && - tileLibBackend == TileLibBackend::PTODSL) - expandOptions = resolveExpandTileOpOptions(argc, argv); if (effectiveBackend == PTOBackend::VPTO && !hasTileOpsToExpand) { if (ptoPrintSeamIR || !ptoSeamIRFile.empty()) { @@ -3389,7 +3213,7 @@ int mlir::pto::compilePTOASModule( return 1; } if (failed(runVPTOBackendPipeline(module, hasTileOpsToExpand, - /*expandOptions=*/nullptr))) + context.getTileLibService()))) return 1; return emitVPTOBackendResult(*module, result, emitVPTOHostStub, context.getCANNVersionOrDefault()); @@ -3427,13 +3251,9 @@ int mlir::pto::compilePTOASModule( // PTODSL legality discovery happens on tile-native PTO IR before fusion. // Fusion may later filter the ordered `candidates` array; ExpandTileOp // consumes the first candidate that remains. - if (!isA2A3 && expandOptions && - expandOptions->tileLibBackend == "ptodsl") { - auto insertOptions = - buildInsertTemplateAttributesOptions(*expandOptions); - pm.addPass( - pto::createInsertTemplateAttributesPass(insertOptions)); - } + if (!isA2A3 && hasTileOpsToExpand) + pm.addPass(pto::createInsertTemplateAttributesPass( + context.getTileLibService())); // Keep frontend fusion on tile-native PTO IR and annotate last_use directly // on scheduled block-local spans before the shared mainline lowers tiles. @@ -3569,12 +3389,6 @@ int mlir::pto::compilePTOASModule( if (ptoPrintSeamIR) printSharedPreBackendSeamIR(*module); - // The PTODSL daemon is needed before the main pipeline for metadata. - // Legacy TileLang can still be resolved lazily immediately before - // ExpandTileOp, preserving the prior --emit-pto-ir behavior. - if (hasTileOpsToExpand && !expandOptions) - expandOptions = resolveExpandTileOpOptions(argc, argv); - if (ptoPrintSeamIR) { module->print(llvm::errs()); llvm::errs() << "\n"; @@ -3583,8 +3397,7 @@ int mlir::pto::compilePTOASModule( return 1; if (failed(runVPTOBackendPipeline( - module, hasTileOpsToExpand, - expandOptions ? &*expandOptions : nullptr))) + module, hasTileOpsToExpand, context.getTileLibService()))) return 1; return emitVPTOBackendResult(*module, result, emitVPTOHostStub, context.getCANNVersionOrDefault()); diff --git a/tools/ptoas/ptoas.h b/tools/ptoas/ptoas.h index cf5b369ef7..b2d317b043 100644 --- a/tools/ptoas/ptoas.h +++ b/tools/ptoas/ptoas.h @@ -13,6 +13,7 @@ #include "PTO/Compiler/CompilerApi.h" #include "PTO/Transforms/VPTOLLVMEmitter.h" #include "VFSIMTSizePatcher.h" +#include "PTO/Transforms/TileLibService.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/MLIRContext.h" #include "mlir/Support/LogicalResult.h" @@ -62,6 +63,9 @@ class PTOASContext { public: PTOASContext(DialectRegistry ®istry, llvm::StringRef outputPath, int argc, char **argv); + PTOASContext(MLIRContext &borrowedContext, + std::shared_ptr tileLibService, + llvm::StringRef outputPath, int argc, char **argv); ~PTOASContext(); LogicalResult initializeEnvironment(bool requiresToolchain, @@ -69,6 +73,7 @@ class PTOASContext { void initializeMLIRContext(); MLIRContext &getMLIRContext(); + std::shared_ptr getTileLibService() const; void setArch(std::string value); llvm::StringRef getArch() const; @@ -94,7 +99,9 @@ class PTOASContext { std::string &path); private: - MLIRContext mlirContext; + std::unique_ptr ownedMlirContext; + MLIRContext *mlirContext = nullptr; + std::shared_ptr tileLibService; std::string outputPath; std::string arch; BackendInfo backendInfo; @@ -135,6 +142,9 @@ void loadPTOASDialects(MLIRContext &context); // Reusable driver entry shared by the Python extension and standalone CLI. PTOAS_COMPILER_EXPORT int runPTOAS(int argc, char **argv); +PTOAS_COMPILER_EXPORT int +runPTOAS(int argc, char **argv, MLIRContext &borrowedContext, + std::shared_ptr tileLibService); // Attach textual-.pto SSA name hints (function args, block args, op results) // to the parsed module's Locations as debug metadata. Called by the driver From c40ae23265e334fc297bc357856b04381394337c Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Mon, 3 Aug 2026 14:38:58 +0800 Subject: [PATCH 063/122] fix(tilelib): restrict metadata discovery to VPTO --- include/PTO/Transforms/TileLibService.h | 4 ++-- ptodsl/ptodsl/tilelib/_compiler_runtime.py | 4 ++-- tools/ptoas/ptoas.cpp | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/include/PTO/Transforms/TileLibService.h b/include/PTO/Transforms/TileLibService.h index c28b19654a..16338d3df2 100644 --- a/include/PTO/Transforms/TileLibService.h +++ b/include/PTO/Transforms/TileLibService.h @@ -1,8 +1,8 @@ // Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms of +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of // CANN Open Software License Agreement Version 2.0 (the "License"). // Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. diff --git a/ptodsl/ptodsl/tilelib/_compiler_runtime.py b/ptodsl/ptodsl/tilelib/_compiler_runtime.py index 23a183273d..0ae7b6afb3 100644 --- a/ptodsl/ptodsl/tilelib/_compiler_runtime.py +++ b/ptodsl/ptodsl/tilelib/_compiler_runtime.py @@ -1,8 +1,8 @@ # Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software, you can redistribute it and/or modify it under the terms of +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of # CANN Open Software License Agreement Version 2.0 (the "License"). # Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 089e079108..4f690e1b35 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -3251,7 +3251,7 @@ int mlir::pto::compilePTOASModule( // PTODSL legality discovery happens on tile-native PTO IR before fusion. // Fusion may later filter the ordered `candidates` array; ExpandTileOp // consumes the first candidate that remains. - if (!isA2A3 && hasTileOpsToExpand) + if (!isA2A3 && effectiveBackend == PTOBackend::VPTO && hasTileOpsToExpand) pm.addPass(pto::createInsertTemplateAttributesPass( context.getTileLibService())); From 4fbe647be8962f6445fd19a3715980e0dcda7047 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Mon, 3 Aug 2026 17:09:43 +0800 Subject: [PATCH 064/122] fix(tilelib): release Python MLIR wrappers after native clone --- ptodsl/ptodsl/tilelib/_compiler_runtime.py | 13 ++++++++++++- tools/ptoas/NativeModule.cpp | 8 ++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/ptodsl/ptodsl/tilelib/_compiler_runtime.py b/ptodsl/ptodsl/tilelib/_compiler_runtime.py index 0ae7b6afb3..a83b7430a8 100644 --- a/ptodsl/ptodsl/tilelib/_compiler_runtime.py +++ b/ptodsl/ptodsl/tilelib/_compiler_runtime.py @@ -69,4 +69,15 @@ def materialize( return module, descriptor.name -__all__ = ["materialize", "metadata"] +def release_module_wrappers(context) -> None: + """Release Python operation wrappers after a native caller clones a module. + + ``Module`` remains the owner of the native module. The binding's live + operation map otherwise retains detached wrappers after the C++ handoff, + so a later materialization can collide with a reused native operation + address. + """ + context._clear_live_operations() + + +__all__ = ["materialize", "metadata", "release_module_wrappers"] diff --git a/tools/ptoas/NativeModule.cpp b/tools/ptoas/NativeModule.cpp index 36b12ce341..3119c83e19 100644 --- a/tools/ptoas/NativeModule.cpp +++ b/tools/ptoas/NativeModule.cpp @@ -76,6 +76,14 @@ class PythonTileLibService final : public mlir::pto::TileLibService { mlir::ModuleOp source = unwrap(rawModule); auto cloned = mlir::cast(source->clone()); + + // Python's MLIR binding keeps operation wrappers in the context's + // live-operation map (and Module.operation keeps the Module alive). + // Clear those non-owning wrappers before the PyModule destructor + // releases its native module. This context is private to this PTOAS + // invocation, so no unrelated Python operation wrappers are invalidated. + getRuntime().attr("release_module_wrappers")(contextOwner); + mlir::pto::TileLibMaterialization materialization{ mlir::OwningOpRef(cloned), py::cast(result[1])}; From 1d2c918ef5a6e76f7d3d476523eaadf72b9271b1 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Mon, 3 Aug 2026 18:08:56 +0800 Subject: [PATCH 065/122] refactor(tilelib): borrow source module during materialization --- include/PTO/Transforms/TileLibService.h | 22 +-- lib/PTO/Transforms/ExpandTileOp.cpp | 162 ++++++++++----------- ptodsl/ptodsl/tilelib/_compiler_runtime.py | 13 +- tools/ptoas/NativeModule.cpp | 19 +-- 4 files changed, 97 insertions(+), 119 deletions(-) diff --git a/include/PTO/Transforms/TileLibService.h b/include/PTO/Transforms/TileLibService.h index 16338d3df2..9dc4bf2859 100644 --- a/include/PTO/Transforms/TileLibService.h +++ b/include/PTO/Transforms/TileLibService.h @@ -12,7 +12,9 @@ #include "mlir/IR/BuiltinOps.h" #include "mlir/Support/LogicalResult.h" -#include +#include "llvm/ADT/STLFunctionalExtras.h" +#include "llvm/ADT/StringRef.h" + #include namespace mlir::pto { @@ -29,14 +31,13 @@ struct TileLibMaterializationRequest { std::string candidateId; }; -struct TileLibMaterialization { - OwningOpRef module; - std::string entrySymbol; -}; +using TileLibMaterializationCallback = + llvm::function_ref; -/// C++ ownership boundary for a TileLib implementation. Implementations return -/// a C++-owned source ModuleOp in the requested context. The caller may then -/// clone/import its generated functions into the caller module. +/// Synchronous handoff for a materialized TileLib implementation. The source +/// module is borrowed and remains owned by the service for the duration of the +/// callback. Consumers must clone/import any operations they need before the +/// callback returns. class TileLibService { public: virtual ~TileLibService() = default; @@ -44,9 +45,10 @@ class TileLibService { virtual FailureOr getMetadata(const TileLibMaterializationRequest &request) = 0; - virtual FailureOr + virtual LogicalResult materialize(const TileLibMaterializationRequest &request, - MLIRContext &context) = 0; + MLIRContext &context, + TileLibMaterializationCallback callback) = 0; }; } // namespace mlir::pto diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index a67acd9599..9a9a733517 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -975,8 +975,8 @@ static std::string buildContextAttrsJson(const SpecKey &key) { // ============================================================================ // Materialize PTODSL in the host Python interpreter and import its functions. -// The service clones the Python-owned source module before returning, so this -// pass only handles C++-owned IR in the current MLIRContext. +// The service borrows the source module only for the synchronous callback; +// this pass clones the required functions into the caller module there. // ============================================================================ func::FuncOp ExpandState::invokeInProcessTileLib(const SpecKey &key, StringRef candidateId, @@ -992,98 +992,96 @@ func::FuncOp ExpandState::invokeInProcessTileLib(const SpecKey &key, request.contextAttrsJson = buildContextAttrsJson(key); request.candidateId = candidateId.str(); - FailureOr materializationOr = - tileLibService->materialize(request, *ctx); - if (failed(materializationOr)) { - llvm::errs() << "ExpandTileOp: in-process PTODSL materialization failed\n"; - return nullptr; - } + func::FuncOp importedEntry; + LogicalResult materializationResult = tileLibService->materialize( + request, *ctx, [&](ModuleOp sourceModule, StringRef entrySymbol) { + if (!sourceModule || sourceModule.getContext() != ctx) { + llvm::errs() << "ExpandTileOp: in-process PTODSL returned a module from " + "a different MLIRContext\n"; + return failure(); + } - pto::TileLibMaterialization materialization = - std::move(*materializationOr); - OwningOpRef sourceModule = std::move(materialization.module); - if (!sourceModule || sourceModule->getContext() != ctx) { - llvm::errs() << "ExpandTileOp: in-process PTODSL returned a module from " - "a different MLIRContext\n"; - return nullptr; - } + auto sourceEntry = sourceModule.lookupSymbol(entrySymbol); + if (!sourceEntry) { + llvm::errs() << "ExpandTileOp: in-process PTODSL entry symbol @" + << entrySymbol << " was not found\n"; + return failure(); + } - auto sourceEntry = sourceModule->lookupSymbol( - materialization.entrySymbol); - if (!sourceEntry) { - llvm::errs() << "ExpandTileOp: in-process PTODSL entry symbol @" - << materialization.entrySymbol << " was not found\n"; - return nullptr; - } + SmallVector sourceFuncs; + for (func::FuncOp fn : sourceModule.getOps()) + sourceFuncs.push_back(fn); + if (sourceFuncs.empty()) { + llvm::errs() << "ExpandTileOp: in-process PTODSL returned no func.func\n"; + return failure(); + } - SmallVector sourceFuncs; - for (func::FuncOp fn : sourceModule->getOps()) - sourceFuncs.push_back(fn); - if (sourceFuncs.empty()) { - llvm::errs() << "ExpandTileOp: in-process PTODSL returned no func.func\n"; - return nullptr; - } + std::string uniqueName = buildUniqueFunctionBaseName(key); + if (!candidateId.empty()) + uniqueName += "__" + candidateId.str(); - std::string uniqueName = buildUniqueFunctionBaseName(key); - if (!candidateId.empty()) - uniqueName += "__" + candidateId.str(); - - SymbolTable targetSymTable(mod); - if (auto existingFunc = targetSymTable.lookup(uniqueName)) - return cast(existingFunc); - - llvm::StringMap plannedSymbols; - for (func::FuncOp fn : sourceFuncs) { - std::string newName = fn == sourceEntry - ? uniqueName - : uniqueName + "__" + std::string(fn.getSymName()); - if (targetSymTable.lookup(newName)) { - llvm::errs() << "ExpandTileOp: imported PTODSL symbol collision at @" - << newName << "\n"; - return nullptr; + SymbolTable targetSymTable(mod); + if (auto existingFunc = targetSymTable.lookup(uniqueName)) { + importedEntry = cast(existingFunc); + return success(); } - plannedSymbols[fn.getSymName()] = std::move(newName); - } - OpBuilder builder(ctx); - builder.setInsertionPointToEnd(mod.getBody()); - SmallVector clonedFuncs; - for (func::FuncOp fn : sourceFuncs) { - IRMapping mapping; - auto cloned = cast(builder.clone(*fn, mapping)); - cloned.setName(plannedSymbols.lookup(fn.getSymName())); - cloned.setVisibility(SymbolTable::Visibility::Private); - clonedFuncs.push_back(cloned); - } + llvm::StringMap plannedSymbols; + for (func::FuncOp fn : sourceFuncs) { + std::string newName = fn == sourceEntry + ? uniqueName + : uniqueName + "__" + std::string(fn.getSymName()); + if (targetSymTable.lookup(newName)) { + llvm::errs() << "ExpandTileOp: imported PTODSL symbol collision at @" + << newName << "\n"; + return failure(); + } + plannedSymbols[fn.getSymName()] = std::move(newName); + } + + OpBuilder builder(ctx); + builder.setInsertionPointToEnd(mod.getBody()); + SmallVector clonedFuncs; + for (func::FuncOp fn : sourceFuncs) { + IRMapping mapping; + auto cloned = cast(builder.clone(*fn, mapping)); + cloned.setName(plannedSymbols.lookup(fn.getSymName())); + cloned.setVisibility(SymbolTable::Visibility::Private); + clonedFuncs.push_back(cloned); + } - for (func::FuncOp fn : clonedFuncs) { - for (const auto &renamed : plannedSymbols) { - if (failed(SymbolTable::replaceAllSymbolUses( - StringAttr::get(ctx, renamed.getKey()), - StringAttr::get(ctx, renamed.getValue()), fn))) { - llvm::errs() << "ExpandTileOp: failed to rewrite imported symbol @" - << renamed.getKey() << " in @" << fn.getSymName() - << "\n"; - for (func::FuncOp imported : clonedFuncs) - imported.erase(); - return nullptr; + for (func::FuncOp fn : clonedFuncs) { + for (const auto &renamed : plannedSymbols) { + if (failed(SymbolTable::replaceAllSymbolUses( + StringAttr::get(ctx, renamed.getKey()), + StringAttr::get(ctx, renamed.getValue()), fn))) { + llvm::errs() << "ExpandTileOp: failed to rewrite imported symbol @" + << renamed.getKey() << " in @" << fn.getSymName() + << "\n"; + for (func::FuncOp imported : clonedFuncs) + imported.erase(); + return failure(); + } } } - } - func::FuncOp entry = mod.lookupSymbol(uniqueName); - if (!entry) { - llvm::errs() << "ExpandTileOp: failed to import PTODSL entry @" - << materialization.entrySymbol << "\n"; + importedEntry = mod.lookupSymbol(uniqueName); + if (!importedEntry) { + llvm::errs() << "ExpandTileOp: failed to import PTODSL entry @" + << entrySymbol << "\n"; + return failure(); + } + if (!importedEntry->hasAttr("pto.tilelang.instance")) + llvm::errs() << "ExpandTileOp: warning: in-process PTODSL entry @" + << importedEntry.getSymName() + << " missing pto.tilelang.instance attribute\n"; + return success(); + }); + if (failed(materializationResult)) { + llvm::errs() << "ExpandTileOp: in-process PTODSL materialization failed\n"; return nullptr; } - if (!entry->hasAttr("pto.tilelang.instance")) { - llvm::errs() << "ExpandTileOp: warning: in-process PTODSL entry @" - << entry.getSymName() - << " missing pto.tilelang.instance attribute\n"; - } - - return entry; + return importedEntry; } // ============================================================================ diff --git a/ptodsl/ptodsl/tilelib/_compiler_runtime.py b/ptodsl/ptodsl/tilelib/_compiler_runtime.py index a83b7430a8..0ae7b6afb3 100644 --- a/ptodsl/ptodsl/tilelib/_compiler_runtime.py +++ b/ptodsl/ptodsl/tilelib/_compiler_runtime.py @@ -69,15 +69,4 @@ def materialize( return module, descriptor.name -def release_module_wrappers(context) -> None: - """Release Python operation wrappers after a native caller clones a module. - - ``Module`` remains the owner of the native module. The binding's live - operation map otherwise retains detached wrappers after the C++ handoff, - so a later materialization can collide with a reused native operation - address. - """ - context._clear_live_operations() - - -__all__ = ["materialize", "metadata", "release_module_wrappers"] +__all__ = ["materialize", "metadata"] diff --git a/tools/ptoas/NativeModule.cpp b/tools/ptoas/NativeModule.cpp index 3119c83e19..ff81286144 100644 --- a/tools/ptoas/NativeModule.cpp +++ b/tools/ptoas/NativeModule.cpp @@ -45,9 +45,10 @@ class PythonTileLibService final : public mlir::pto::TileLibService { } } - mlir::FailureOr + mlir::LogicalResult materialize(const mlir::pto::TileLibMaterializationRequest &request, - mlir::MLIRContext &context) override { + mlir::MLIRContext &context, + mlir::pto::TileLibMaterializationCallback callback) override { py::gil_scoped_acquire acquire; try { MlirContext pythonContext = py::cast(contextOwner); @@ -75,19 +76,7 @@ class PythonTileLibService final : public mlir::pto::TileLibService { } mlir::ModuleOp source = unwrap(rawModule); - auto cloned = mlir::cast(source->clone()); - - // Python's MLIR binding keeps operation wrappers in the context's - // live-operation map (and Module.operation keeps the Module alive). - // Clear those non-owning wrappers before the PyModule destructor - // releases its native module. This context is private to this PTOAS - // invocation, so no unrelated Python operation wrappers are invalidated. - getRuntime().attr("release_module_wrappers")(contextOwner); - - mlir::pto::TileLibMaterialization materialization{ - mlir::OwningOpRef(cloned), - py::cast(result[1])}; - return materialization; + return callback(source, py::cast(result[1])); } catch (const py::error_already_set &error) { llvm::errs() << "TileLib: PTODSL materialization raised Python " "exception:\n" From 2fb1e8a5c6259a17e4c2651cefe707096404e108 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Mon, 3 Aug 2026 19:07:18 +0800 Subject: [PATCH 066/122] fix(tilelib): invalidate wrappers after borrowed import --- ptodsl/ptodsl/tilelib/_compiler_runtime.py | 17 ++++++++++++++++- tools/ptoas/NativeModule.cpp | 10 +++++++++- 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/ptodsl/ptodsl/tilelib/_compiler_runtime.py b/ptodsl/ptodsl/tilelib/_compiler_runtime.py index 0ae7b6afb3..c56b69f257 100644 --- a/ptodsl/ptodsl/tilelib/_compiler_runtime.py +++ b/ptodsl/ptodsl/tilelib/_compiler_runtime.py @@ -69,4 +69,19 @@ def materialize( return module, descriptor.name -__all__ = ["materialize", "metadata"] +def invalidate_materialized_module_wrappers(context) -> None: + """Invalidate wrappers after native code imports a materialized module. + + PTOAS uses a dedicated Python MLIR context for each compiler invocation. + Native code only borrows the source module during a synchronous callback, + but the binding's live-operation map otherwise retains wrappers after that + callback. No Python operation from this materialization remains usable. + """ + context._clear_live_operations() + + +__all__ = [ + "invalidate_materialized_module_wrappers", + "materialize", + "metadata", +] diff --git a/tools/ptoas/NativeModule.cpp b/tools/ptoas/NativeModule.cpp index ff81286144..427983ed39 100644 --- a/tools/ptoas/NativeModule.cpp +++ b/tools/ptoas/NativeModule.cpp @@ -76,7 +76,15 @@ class PythonTileLibService final : public mlir::pto::TileLibService { } mlir::ModuleOp source = unwrap(rawModule); - return callback(source, py::cast(result[1])); + mlir::LogicalResult callbackResult = + callback(source, py::cast(result[1])); + + // The source module remains Python-owned, but native import is the last + // permitted use of its operation wrappers. Keep the binding's live map + // synchronized before the next materialization creates detached ops. + getRuntime().attr("invalidate_materialized_module_wrappers")( + contextOwner); + return callbackResult; } catch (const py::error_already_set &error) { llvm::errs() << "TileLib: PTODSL materialization raised Python " "exception:\n" From 05ee0f84587207df3f613cd9d83b01c121712ded Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Mon, 3 Aug 2026 19:59:56 +0800 Subject: [PATCH 067/122] fix(ptodsl): break tile valid-shape wrapper cycles --- ptodsl/ptodsl/_surface_values.py | 3 ++- ptodsl/ptodsl/tilelib/_compiler_runtime.py | 17 +---------------- ptodsl/tests/test_tilelib_render.py | 13 +++++++++++++ tools/ptoas/NativeModule.cpp | 10 +--------- 4 files changed, 17 insertions(+), 26 deletions(-) diff --git a/ptodsl/ptodsl/_surface_values.py b/ptodsl/ptodsl/_surface_values.py index 0147e9bda7..dff4143395 100644 --- a/ptodsl/ptodsl/_surface_values.py +++ b/ptodsl/ptodsl/_surface_values.py @@ -10,6 +10,7 @@ from __future__ import annotations import re +import weakref from dataclasses import dataclass from ._diagnostics import native_python_control_flow_error @@ -519,7 +520,7 @@ class _TileValidShapeView: """Tuple-like proxy that lowers `tile.valid_shape[i]` on demand.""" def __init__(self, tile: "TileValue"): - self._tile = tile + self._tile = weakref.proxy(tile) self._cache: dict[int, object] = {} def __getitem__(self, index: int): diff --git a/ptodsl/ptodsl/tilelib/_compiler_runtime.py b/ptodsl/ptodsl/tilelib/_compiler_runtime.py index c56b69f257..0ae7b6afb3 100644 --- a/ptodsl/ptodsl/tilelib/_compiler_runtime.py +++ b/ptodsl/ptodsl/tilelib/_compiler_runtime.py @@ -69,19 +69,4 @@ def materialize( return module, descriptor.name -def invalidate_materialized_module_wrappers(context) -> None: - """Invalidate wrappers after native code imports a materialized module. - - PTOAS uses a dedicated Python MLIR context for each compiler invocation. - Native code only borrows the source module during a synchronous callback, - but the binding's live-operation map otherwise retains wrappers after that - callback. No Python operation from this materialization remains usable. - """ - context._clear_live_operations() - - -__all__ = [ - "invalidate_materialized_module_wrappers", - "materialize", - "metadata", -] +__all__ = ["materialize", "metadata"] diff --git a/ptodsl/tests/test_tilelib_render.py b/ptodsl/tests/test_tilelib_render.py index 0f07d4668f..ac6821c130 100644 --- a/ptodsl/tests/test_tilelib_render.py +++ b/ptodsl/tests/test_tilelib_render.py @@ -83,6 +83,19 @@ def test_materialize_uses_borrowed_context_and_returns_fresh_modules(self): self.assertTrue(second.operation.verify()) self.assertIn("func.func @template_tadd", artifact.mlir_text()) + def test_materialized_surface_wrappers_release_without_cycle_collection(self): + context = Context() + pto_dialect.register_dialect(context, load=True) + spec = TileSpec(shape=(8, 64), dtype=f32) + + module = template_tadd.specialize( + src0=spec, src1=spec, dst=spec + ).materialize(context) + self.assertEqual(context._get_live_operation_count(), 0) + + del module + self.assertEqual(context._get_live_operation_count(), 0) + if __name__ == "__main__": unittest.main() diff --git a/tools/ptoas/NativeModule.cpp b/tools/ptoas/NativeModule.cpp index 427983ed39..ff81286144 100644 --- a/tools/ptoas/NativeModule.cpp +++ b/tools/ptoas/NativeModule.cpp @@ -76,15 +76,7 @@ class PythonTileLibService final : public mlir::pto::TileLibService { } mlir::ModuleOp source = unwrap(rawModule); - mlir::LogicalResult callbackResult = - callback(source, py::cast(result[1])); - - // The source module remains Python-owned, but native import is the last - // permitted use of its operation wrappers. Keep the binding's live map - // synchronized before the next materialization creates detached ops. - getRuntime().attr("invalidate_materialized_module_wrappers")( - contextOwner); - return callbackResult; + return callback(source, py::cast(result[1])); } catch (const py::error_already_set &error) { llvm::errs() << "TileLib: PTODSL materialization raised Python " "exception:\n" From b42b34079012e73557de891b9c37cb7544a825fd Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Mon, 3 Aug 2026 20:40:43 +0800 Subject: [PATCH 068/122] refactor(ptodsl): make valid-shape views ephemeral --- ptodsl/ptodsl/_surface_values.py | 41 ++++++++++++------------ ptodsl/ptodsl/tilelib/_render_runtime.py | 2 +- 2 files changed, 22 insertions(+), 21 deletions(-) diff --git a/ptodsl/ptodsl/_surface_values.py b/ptodsl/ptodsl/_surface_values.py index dff4143395..375806b615 100644 --- a/ptodsl/ptodsl/_surface_values.py +++ b/ptodsl/ptodsl/_surface_values.py @@ -10,7 +10,6 @@ from __future__ import annotations import re -import weakref from dataclasses import dataclass from ._diagnostics import native_python_control_flow_error @@ -520,35 +519,41 @@ class _TileValidShapeView: """Tuple-like proxy that lowers `tile.valid_shape[i]` on demand.""" def __init__(self, tile: "TileValue"): - self._tile = weakref.proxy(tile) - self._cache: dict[int, object] = {} + self._tile = tile def __getitem__(self, index: int): - logical_rank = len(self._tile.shape) if self._tile.shape is not None else 2 + return self._tile._get_valid_shape_dim(index) + + +class TileValue(_SurfaceValue, Tile): + """Author-facing tile handle with surface-style accessors.""" + + def _get_valid_shape_dim(self, index: int): + logical_rank = len(self.shape) if self.shape is not None else 2 allowed = {0} if logical_rank == 1 else {0, 1} if index not in allowed: if logical_rank == 1: raise IndexError("PTODSL rank-1 tile.valid_shape currently supports only index 0") raise IndexError("PTODSL tile.valid_shape currently supports indices 0 and 1") - cached = self._cache.get(index) + cached = self._valid_shape_cache.get(index) if cached is not None: return cached - if self._tile.static_valid_shape is not None: - dim = self._tile.static_valid_shape[index] + if self.static_valid_shape is not None: + dim = self.static_valid_shape[index] if dim is not None: value = _index_const(dim) if _is_python_index_literal(dim) else unwrap_surface_value(dim) value = wrap_surface_value(value) - self._cache[index] = value + self._valid_shape_cache[index] = value return value try: if logical_rank == 1: - value = wrap_surface_value(_pto.TileValidColsOp(self._tile.value).result) + value = wrap_surface_value(_pto.TileValidColsOp(self.value).result) elif index == 0: - value = wrap_surface_value(_pto.TileValidRowsOp(self._tile.value).result) + value = wrap_surface_value(_pto.TileValidRowsOp(self.value).result) else: - value = wrap_surface_value(_pto.TileValidColsOp(self._tile.value).result) + value = wrap_surface_value(_pto.TileValidColsOp(self.value).result) except Exception: - static_dim = _fallback_static_valid_dim(self._tile.type, index) + static_dim = _fallback_static_valid_dim(self.type, index) if static_dim is None: raise RuntimeError( "tile.valid_shape could not be lowered because the current " @@ -556,13 +561,9 @@ def __getitem__(self, index: int): "the tile type does not carry a recoverable static bound" ) from None value = wrap_surface_value(_index_const(static_dim)) - self._cache[index] = value + self._valid_shape_cache[index] = value return value - -class TileValue(_SurfaceValue, Tile): - """Author-facing tile handle with surface-style accessors.""" - def __init__( self, value, @@ -592,11 +593,11 @@ def __init__( self.static_valid_shape = tuple(valid_shape) if valid_shape is not None else ( parsed["valid_dims"] if parsed is not None else None ) - self._valid_shape = _TileValidShapeView(self) + self._valid_shape_cache: dict[int, object] = {} @property def valid_shape(self): - return self._valid_shape + return _TileValidShapeView(self) @valid_shape.setter def valid_shape(self, dims): @@ -604,7 +605,7 @@ def valid_shape(self, dims): set_tile_valid_shape(self, dims) self.static_valid_shape = tuple(dims) - self._valid_shape._cache.clear() + self._valid_shape_cache.clear() @property def surface_metadata(self): diff --git a/ptodsl/ptodsl/tilelib/_render_runtime.py b/ptodsl/ptodsl/tilelib/_render_runtime.py index b79386bd81..4ecbca900d 100644 --- a/ptodsl/ptodsl/tilelib/_render_runtime.py +++ b/ptodsl/ptodsl/tilelib/_render_runtime.py @@ -69,7 +69,7 @@ def __init__(self, value, spec: TileSpec): ) # Force the dynamic valid-shape ops to match the tilelang render. self.static_valid_shape = None - self._valid_shape._cache.clear() + self._valid_shape_cache.clear() self._template_static_valid_shape = tuple(spec.valid_shape or spec.shape) self._template_config = _TemplateTileConfig( b_layout=spec.b_layout, From 78deb108dd9f622f3ceb0a0f1a00dea09aea2ea1 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Tue, 4 Aug 2026 01:52:42 +0800 Subject: [PATCH 069/122] refactor(tilelib): make service passes session-only --- ...todsl-tilelib-template-selection-design.md | 4 + include/PTO/Transforms/CMakeLists.txt | 3 + include/PTO/Transforms/Passes.h | 2 - include/PTO/Transforms/Passes.td | 43 --------- include/PTO/Transforms/TileLibPasses.td | 62 ++++++++++++ lib/PTO/Transforms/ExpandTileOp.cpp | 8 +- .../Transforms/InsertTemplateAttributes.cpp | 6 +- ptodsl/ptodsl/tilelib/_compiler_runtime.py | 9 +- ptodsl/ptodsl/tilelib/decorator.py | 14 --- ptodsl/tests/test_tilelib_render.py | 38 +++++--- .../script/run_a5_st_all_parallel.py | 96 +++++++++---------- tools/ptoas/ptoas.cpp | 1 - 12 files changed, 149 insertions(+), 137 deletions(-) create mode 100644 include/PTO/Transforms/TileLibPasses.td diff --git a/docs/designs/ptodsl-tilelib-template-selection-design.md b/docs/designs/ptodsl-tilelib-template-selection-design.md index f5055c1e53..34b2914817 100644 --- a/docs/designs/ptodsl-tilelib-template-selection-design.md +++ b/docs/designs/ptodsl-tilelib-template-selection-design.md @@ -66,6 +66,10 @@ before later passes can make candidate information harder to reconstruct. `ExpandTileOp` still renders from the current MLIR operands so the helper body matches the actual operand types and view metadata that survived to expansion. +Both stages are compiler-session passes. They require the `TileLibService` +owned by `PTOASContext`, are constructed explicitly by the PTOAS pipeline, and +are not registered as standalone textual passes. + ## Template Metadata PTODSL template authors register versions through `tilelib.tile_template`. diff --git a/include/PTO/Transforms/CMakeLists.txt b/include/PTO/Transforms/CMakeLists.txt index 5cf85fabb6..16fa1141d7 100644 --- a/include/PTO/Transforms/CMakeLists.txt +++ b/include/PTO/Transforms/CMakeLists.txt @@ -14,5 +14,8 @@ set(LLVM_TARGET_DEFINITIONS Passes.td) mlir_tablegen(Passes.h.inc -gen-pass-decls -name PTO) +set(LLVM_TARGET_DEFINITIONS TileLibPasses.td) +mlir_tablegen(TileLibPasses.h.inc -gen-pass-decls -name PTOTileLib) + # [关键] 重命名为 PTOPassesIncGen (去掉 MLIR 前缀) add_public_tablegen_target(PTOPassesIncGen) diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index 50d9fdfd20..624446d5cf 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -133,10 +133,8 @@ std::unique_ptr createVMILegalizeArithSelectPass(); std::unique_ptr createVMILowerUnifiedToLegacyPass(); std::unique_ptr createVMINormalizeSignlessIntToUnsignedPass(); std::unique_ptr createVMIToVPTOPass(); -std::unique_ptr createInsertTemplateAttributesPass(); std::unique_ptr createInsertTemplateAttributesPass( std::shared_ptr tileLibService); -std::unique_ptr createExpandTileOpPass(); std::unique_ptr createExpandTileOpPass(std::shared_ptr tileLibService); std::unique_ptr createFoldTileBufIntrinsicsPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index a21ec988e2..ab9df2ff97 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -529,49 +529,6 @@ def PTOResolveReservedBuffers : Pass<"pto-resolve-reserved-buffers", "ModuleOp"> ]; } -def InsertTemplateAttributes - : Pass<"pto-insert-template-attributes", "ModuleOp"> { - let summary = "Attach legal PTODSL template candidates to tile operations"; - let description = [{ - Queries the compiler's in-process PTODSL TileLib service for legal template - candidates and stores the compact candidate list on each tile operation as - the `candidates` attribute. Each candidate contains only id, name, - loop_depth, postupdate, and tail metadata. - }]; - let constructor = "mlir::pto::createInsertTemplateAttributesPass()"; - let dependentDialects = [ - "mlir::pto::PTODialect", - "mlir::func::FuncDialect" - ]; -} - -def ExpandTileOp : Pass<"pto-expand-tile-op", "ModuleOp"> { - let summary = "Expand tile ops into calls to TileLib template functions"; - let description = [{ - Expands tile-level operations (pto.tadd, pto.tsub, etc.) by asking the - compiler's in-process PTODSL TileLib service to instantiate template - libraries in the current MLIRContext. The generated template functions use - tile_buf parameters and contain vector-level implementations (pto.vecscope, - pto.vlds, pto.vadd, pto.vsts, etc.). - - Each tile op is replaced by a func.call to the generated template function, - with tile_buf operands passed directly (no type bridging). - - After this pass, the Inline pass inlines template bodies, and - FoldTileBufIntrinsics resolves tile_buf_addr / tile_valid_rows / - tile_valid_cols. - }]; - let constructor = "mlir::pto::createExpandTileOpPass()"; - let dependentDialects = [ - "mlir::pto::PTODialect", - "mlir::memref::MemRefDialect", - "mlir::arith::ArithDialect", - "mlir::func::FuncDialect", - "mlir::scf::SCFDialect", - "mlir::vector::VectorDialect" - ]; -} - def FoldTileBufIntrinsics : Pass<"pto-fold-tile-buf-intrinsics", "mlir::func::FuncOp"> { let summary = "Fold structured-view intrinsics after template inlining"; let description = [{ diff --git a/include/PTO/Transforms/TileLibPasses.td b/include/PTO/Transforms/TileLibPasses.td new file mode 100644 index 0000000000..bb2cfd095a --- /dev/null +++ b/include/PTO/Transforms/TileLibPasses.td @@ -0,0 +1,62 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// TileLib passes require a live compiler-session service and are intentionally +// excluded from the global textual pass registry. Their constructor strings +// suppress generated default factories; PTOAS creates them through the +// service-injecting factories declared in Passes.h. + +#ifndef MLIR_DIALECT_PTO_TILELIB_PASSES +#define MLIR_DIALECT_PTO_TILELIB_PASSES + +include "mlir/Pass/PassBase.td" + +def InsertTemplateAttributes + : Pass<"pto-insert-template-attributes", "ModuleOp"> { + let summary = "Attach legal PTODSL template candidates to tile operations"; + let description = [{ + Queries the compiler's in-process PTODSL TileLib service for legal template + candidates and stores the compact candidate list on each tile operation as + the `candidates` attribute. Each candidate contains only id, name, + loop_depth, postupdate, and tail metadata. + }]; + let constructor = "mlir::pto::createInsertTemplateAttributesPass()"; + let dependentDialects = [ + "mlir::pto::PTODialect", + "mlir::func::FuncDialect" + ]; +} + +def ExpandTileOp : Pass<"pto-expand-tile-op", "ModuleOp"> { + let summary = "Expand tile ops into calls to TileLib template functions"; + let description = [{ + Expands tile-level operations (pto.tadd, pto.tsub, etc.) by asking the + compiler's in-process PTODSL TileLib service to instantiate template + libraries in the current MLIRContext. The generated template functions use + tile_buf parameters and contain vector-level implementations (pto.vecscope, + pto.vlds, pto.vadd, pto.vsts, etc.). + + Each tile op is replaced by a func.call to the generated template function, + with tile_buf operands passed directly (no type bridging). + + After this pass, the Inline pass inlines template bodies, and + FoldTileBufIntrinsics resolves tile_buf_addr / tile_valid_rows / + tile_valid_cols. + }]; + let constructor = "mlir::pto::createExpandTileOpPass()"; + let dependentDialects = [ + "mlir::pto::PTODialect", + "mlir::memref::MemRefDialect", + "mlir::arith::ArithDialect", + "mlir::func::FuncDialect", + "mlir::scf::SCFDialect", + "mlir::vector::VectorDialect" + ]; +} + +#endif // MLIR_DIALECT_PTO_TILELIB_PASSES diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 9a9a733517..3c8530a484 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -63,8 +63,8 @@ namespace mlir { namespace pto { namespace func = ::mlir::func; - #define GEN_PASS_DEF_EXPANDTILEOP - #include "PTO/Transforms/Passes.h.inc" +#define GEN_PASS_DEF_EXPANDTILEOP +#include "PTO/Transforms/TileLibPasses.h.inc" } // namespace pto } // namespace mlir @@ -1210,10 +1210,6 @@ void ExpandTileOpPass::runOnOperation() { namespace mlir { namespace pto { -std::unique_ptr createExpandTileOpPass() { - return std::make_unique(); -} - std::unique_ptr createExpandTileOpPass( std::shared_ptr tileLibService) { return std::make_unique(std::move(tileLibService)); diff --git a/lib/PTO/Transforms/InsertTemplateAttributes.cpp b/lib/PTO/Transforms/InsertTemplateAttributes.cpp index cefd2aeb2f..1ed474429d 100644 --- a/lib/PTO/Transforms/InsertTemplateAttributes.cpp +++ b/lib/PTO/Transforms/InsertTemplateAttributes.cpp @@ -36,7 +36,7 @@ using namespace mlir; namespace mlir { namespace pto { #define GEN_PASS_DEF_INSERTTEMPLATEATTRIBUTES -#include "PTO/Transforms/Passes.h.inc" +#include "PTO/Transforms/TileLibPasses.h.inc" } // namespace pto } // namespace mlir @@ -910,10 +910,6 @@ struct InsertTemplateAttributesPass namespace mlir { namespace pto { -std::unique_ptr createInsertTemplateAttributesPass() { - return std::make_unique(); -} - std::unique_ptr createInsertTemplateAttributesPass( std::shared_ptr tileLibService) { return std::make_unique( diff --git a/ptodsl/ptodsl/tilelib/_compiler_runtime.py b/ptodsl/ptodsl/tilelib/_compiler_runtime.py index 0ae7b6afb3..bc21d1c395 100644 --- a/ptodsl/ptodsl/tilelib/_compiler_runtime.py +++ b/ptodsl/ptodsl/tilelib/_compiler_runtime.py @@ -12,6 +12,7 @@ import json +from ._render_runtime import _TemplateTrace from ._selection import _select_descriptor_and_specs, metadata_request @@ -60,11 +61,11 @@ def materialize( context_attrs, candidate_id or None, ) - artifact = descriptor.specialize( + module = _TemplateTrace( + descriptor, + tile_specs, context_attrs=context_attrs, - **tile_specs, - ) - module = artifact.materialize(context) + ).build_module(context=context) module.operation.verify() return module, descriptor.name diff --git a/ptodsl/ptodsl/tilelib/decorator.py b/ptodsl/ptodsl/tilelib/decorator.py index 0cf4294353..270b3a38a7 100644 --- a/ptodsl/ptodsl/tilelib/decorator.py +++ b/ptodsl/ptodsl/tilelib/decorator.py @@ -56,20 +56,6 @@ def __init__(self, descriptor: TileTemplate, tile_specs: dict, context_attrs=Non self.tile_specs = tile_specs self.context_attrs = dict(context_attrs or {}) - def materialize(self, context): - """Build a fresh source module in the caller-provided context. - - The returned module remains Python-owned. Native callers must keep this - object alive until they have cloned/imported its generated functions. - This bypasses the context-bound ModuleArtifact cache. - """ - return _TemplateTrace( - self.descriptor, - self.tile_specs, - context_attrs=self.context_attrs, - ).build_module(context=context) - - def tile_template(*, op, target="a5", name=None, dtypes=(), layouts=(), memory_spaces=(), constraints=(), priority=0, fusible=False, loop_depth=None, id=None, Tail=None, is_post_update=False, diff --git a/ptodsl/tests/test_tilelib_render.py b/ptodsl/tests/test_tilelib_render.py index ac6821c130..bc69124e03 100644 --- a/ptodsl/tests/test_tilelib_render.py +++ b/ptodsl/tests/test_tilelib_render.py @@ -12,11 +12,13 @@ (ptodsl differs in SSA naming, constant hoisting, index-vs-i32 carry, ptr typing). """ +import json import unittest from pathlib import Path from ptoas.mlir.dialects import pto as pto_dialect from ptoas.mlir.ir import Context +from ptodsl.tilelib._compiler_runtime import materialize from ptodsl.tilelib import TileSpec, f32 from TileOps.a5.tadd import template_tadd @@ -46,6 +48,25 @@ def _render(): return template_tadd.specialize(src0=spec, src1=spec, dst=spec).mlir_text() +def _materialize(context): + tile_spec = { + "kind": "tile", + "shape": [8, 64], + "valid_shape": [8, 64], + "dtype": "f32", + "memory_space": "ub", + "config": {"b_layout": "row_major", "s_layout": "none_box"}, + } + return materialize( + "a5", + "pto.tadd", + json.dumps([tile_spec, tile_spec, tile_spec]), + "{}", + "template_tadd", + context, + ) + + class TileLibRenderTest(unittest.TestCase): def test_renders_structured_abstraction(self): text = _render() @@ -70,27 +91,22 @@ def test_golden_fixture_uses_same_abstraction(self): def test_materialize_uses_borrowed_context_and_returns_fresh_modules(self): context = Context() pto_dialect.register_dialect(context, load=True) - spec = TileSpec(shape=(8, 64), dtype=f32) - artifact = template_tadd.specialize(src0=spec, src1=spec, dst=spec) - - first = artifact.materialize(context) - second = artifact.materialize(context) + first, first_entry = _materialize(context) + second, second_entry = _materialize(context) self.assertIs(first.context, context) self.assertIs(second.context, context) self.assertIsNot(first, second) + self.assertEqual(first_entry, "template_tadd") + self.assertEqual(second_entry, "template_tadd") self.assertTrue(first.operation.verify()) self.assertTrue(second.operation.verify()) - self.assertIn("func.func @template_tadd", artifact.mlir_text()) + self.assertIn("func.func @template_tadd", str(first)) def test_materialized_surface_wrappers_release_without_cycle_collection(self): context = Context() pto_dialect.register_dialect(context, load=True) - spec = TileSpec(shape=(8, 64), dtype=f32) - - module = template_tadd.specialize( - src0=spec, src1=spec, dst=spec - ).materialize(context) + module, _ = _materialize(context) self.assertEqual(context._get_live_operation_count(), 0) del module diff --git a/test/tilelang_st/script/run_a5_st_all_parallel.py b/test/tilelang_st/script/run_a5_st_all_parallel.py index 962636df1e..d4fc1cd4fb 100755 --- a/test/tilelang_st/script/run_a5_st_all_parallel.py +++ b/test/tilelang_st/script/run_a5_st_all_parallel.py @@ -137,60 +137,54 @@ def _run_one(job, args, ptoas_bin, output_root, base_env): "build_dir": str(build_dir), } - try: - with log_path.open("w", encoding="utf-8") as log_handle: - log_handle.write(f"# kind: {kind}\n") - log_handle.write(f"# testcase: {testcase}\n") - log_handle.write(f"# source: {job['target_dir']}\n") - log_handle.write(f"# build: {build_dir}\n") - log_handle.write(f"# PTODSL_CACHE_DIR={env['PTODSL_CACHE_DIR']}\n") - log_handle.write("\n") - - cmake_cmd = [ - "cmake", - "-S", - job["target_dir"], - "-B", - build_dir, - f"-DRUN_MODE={args.run_mode}", - f"-DSOC_VERSION={DEFAULT_SOC_VERSION}", - f"-DTEST_CASE={testcase}", - f"-DPTOAS_BIN={ptoas_bin}", - ] - - rc = _run_logged(cmake_cmd, log_handle, output_root, env) - if rc == 0: - rc = _run_logged( - ["cmake", "--build", build_dir, "--parallel", str(args.build_jobs)], - log_handle, - output_root, - env, - ) + with log_path.open("w", encoding="utf-8") as log_handle: + log_handle.write(f"# kind: {kind}\n") + log_handle.write(f"# testcase: {testcase}\n") + log_handle.write(f"# source: {job['target_dir']}\n") + log_handle.write(f"# build: {build_dir}\n") + log_handle.write(f"# PTODSL_CACHE_DIR={env['PTODSL_CACHE_DIR']}\n") + log_handle.write("\n") + + cmake_cmd = [ + "cmake", + "-S", + job["target_dir"], + "-B", + build_dir, + f"-DRUN_MODE={args.run_mode}", + f"-DSOC_VERSION={DEFAULT_SOC_VERSION}", + f"-DTEST_CASE={testcase}", + f"-DPTOAS_BIN={ptoas_bin}", + ] + + rc = _run_logged(cmake_cmd, log_handle, output_root, env) + if rc == 0: + rc = _run_logged( + ["cmake", "--build", build_dir, "--parallel", str(args.build_jobs)], + log_handle, + output_root, + env, + ) + if rc != 0: + result["returncode"] = rc + result["phase"] = "build" + result["seconds"] = time.time() - started + return result + + case_work_dir = build_dir / "testcase" / testcase + _copy_case_scripts(job["testcase_root"], testcase, case_work_dir) + + for phase, command in ( + ("gen_data", [sys.executable, "gen_data.py"]), + ("run", [build_dir / "bin" / testcase]), + ("compare", [sys.executable, "compare.py"]), + ): + rc = _run_logged(command, log_handle, case_work_dir, env) if rc != 0: result["returncode"] = rc - result["phase"] = "build" - result["seconds"] = time.time() - started - return result - - case_work_dir = build_dir / "testcase" / testcase - _copy_case_scripts(job["testcase_root"], testcase, case_work_dir) - - for phase, command in ( - ("gen_data", [sys.executable, "gen_data.py"]), - ("run", [build_dir / "bin" / testcase]), - ("compare", [sys.executable, "compare.py"]), - ): - rc = _run_logged(command, log_handle, case_work_dir, env) - if rc != 0: - result["returncode"] = rc - result["phase"] = phase - break - finally: - try: - socket_path.unlink(missing_ok=True) - except OSError: - pass + result["phase"] = phase + break result["seconds"] = time.time() - started return result diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 4f690e1b35..fd00bbce2c 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -213,7 +213,6 @@ void mlir::pto::registerPTOASPassesAndCLOptions() { mlir::pto::registerPTOPasses(); mlir::pto::registerPTOInlineLibCall(); mlir::pto::registerFoldTileBufIntrinsics(); - mlir::pto::registerExpandTileOp(); mlir::pto::registerLowerPTOToUBufOps(); mlir::registerPassManagerCLOptions(); } From 1eacbd267c4f953f4dd2da7cd8833d132ca0b0e9 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Tue, 4 Aug 2026 10:03:13 +0800 Subject: [PATCH 070/122] test(tilelib): print session pass IR without registration --- test/lit/vpto/expand_tile_op_ptodsl_tadd.pto | 6 ++++-- test/lit/vpto/vpto_pipeline_vmi_after_tileop_expand.pto | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto b/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto index b7f22c355f..5dd33b1c85 100644 --- a/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto +++ b/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto @@ -13,11 +13,12 @@ // Running without --enable-op-fusion proves metadata insertion is not gated // by fusion. Printing before FusionPlan proves its position when fusion is on. // -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after-all %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --pto-level=level2 --enable-op-fusion --emit-pto-ir --mlir-print-ir-before=pto-fusion-plan %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=PREFUSION -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after-all %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>/dev/null | FileCheck %s --check-prefix=EXPAND +// META: IR Dump After InsertTemplateAttributes // META: pto.tadd // META-SAME: candidates = [ // META-SAME: id = 0 : i64 @@ -33,6 +34,7 @@ // PREFUSION: pto.tadd // PREFUSION-SAME: candidates = [ +// SELECT: IR Dump After ExpandTileOp // SELECT: func.func {{.*}}@{{.*}}__template_tadd // EXPAND: func.func @TADD diff --git a/test/lit/vpto/vpto_pipeline_vmi_after_tileop_expand.pto b/test/lit/vpto/vpto_pipeline_vmi_after_tileop_expand.pto index a87c473260..92ecb930b1 100644 --- a/test/lit/vpto/vpto_pipeline_vmi_after_tileop_expand.pto +++ b/test/lit/vpto/vpto_pipeline_vmi_after_tileop_expand.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o /dev/null --mlir-print-ir-after=pto-expand-tile-op --mlir-print-ir-after=vmi-to-vpto --mlir-print-ir-after=vpto-ptr-normalize --mlir-print-ir-after=pto-infer-vpto-vecscope --mlir-print-ir-after=loop-invariant-code-motion 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o /dev/null --mlir-print-ir-after-all 2>&1 | FileCheck %s // CHECK: IR Dump After ExpandTileOp // CHECK: IR Dump After VMIToVPTO From bac713e22d2c9d39350e8d1c27f214cbac58d167 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Tue, 4 Aug 2026 13:19:47 +0800 Subject: [PATCH 071/122] refactor(tilelib): share process runtime across contexts --- ...todsl-tilelib-template-selection-design.md | 16 ++++- include/PTO/Transforms/CMakeLists.txt | 3 - include/PTO/Transforms/Passes.h | 6 +- include/PTO/Transforms/Passes.td | 43 ++++++++++++ include/PTO/Transforms/TileLibPasses.td | 62 ------------------ include/PTO/Transforms/TileLibService.h | 12 ++++ lib/PTO/Transforms/CMakeLists.txt | 1 + lib/PTO/Transforms/ExpandTileOp.cpp | 21 ++---- .../Transforms/InsertTemplateAttributes.cpp | 22 ++----- lib/PTO/Transforms/TileLibService.cpp | 43 ++++++++++++ ptodsl/ptodsl/tilelib/_compiler_runtime.py | 2 +- ptodsl/ptodsl/tilelib/_render_runtime.py | 16 ++++- ptodsl/ptodsl/tilelib/decorator.py | 2 +- ptodsl/tests/test_ptoas_runtime.py | 51 +++++++++++++++ test/lit/vpto/expand_tile_op_ptodsl_tadd.pto | 6 +- .../vpto_pipeline_vmi_after_tileop_expand.pto | 2 +- tools/ptoas/NativeModule.cpp | 65 ++++++++++++++----- tools/ptoas/driver.cpp | 31 ++++----- tools/ptoas/ptoas.cpp | 22 ++----- tools/ptoas/ptoas.h | 11 +--- 20 files changed, 264 insertions(+), 173 deletions(-) delete mode 100644 include/PTO/Transforms/TileLibPasses.td create mode 100644 lib/PTO/Transforms/TileLibService.cpp create mode 100644 ptodsl/tests/test_ptoas_runtime.py diff --git a/docs/designs/ptodsl-tilelib-template-selection-design.md b/docs/designs/ptodsl-tilelib-template-selection-design.md index 34b2914817..35d795d651 100644 --- a/docs/designs/ptodsl-tilelib-template-selection-design.md +++ b/docs/designs/ptodsl-tilelib-template-selection-design.md @@ -66,9 +66,19 @@ before later passes can make candidate information harder to reconstruct. `ExpandTileOp` still renders from the current MLIR operands so the helper body matches the actual operand types and view metadata that survived to expansion. -Both stages are compiler-session passes. They require the `TileLibService` -owned by `PTOASContext`, are constructed explicitly by the PTOAS pipeline, and -are not registered as standalone textual passes. +Both stages are ordinary registered MLIR passes. They are default-constructible +and obtain the current `MLIRContext` from the operation being transformed. A +process-wide `TileLibRuntime` provides the host `TileLibService`; it owns no +compilation context and receives the current context explicitly for every +materialization. `PTOASContext` continues to own or borrow the context for one +compilation session, so different invocations may use different contexts while +sharing one Python runtime. + +The Python entry keeps the corresponding Python `Context` owner alive for the +complete native compilation call. Compiler materialization requires that +explicit context and never falls back to creating another one. This preserves +normal pass registration, textual pipelines, targeted IR printing, cloning, +and reproducer behavior without storing Python objects in pass instances. ## Template Metadata diff --git a/include/PTO/Transforms/CMakeLists.txt b/include/PTO/Transforms/CMakeLists.txt index 16fa1141d7..5cf85fabb6 100644 --- a/include/PTO/Transforms/CMakeLists.txt +++ b/include/PTO/Transforms/CMakeLists.txt @@ -14,8 +14,5 @@ set(LLVM_TARGET_DEFINITIONS Passes.td) mlir_tablegen(Passes.h.inc -gen-pass-decls -name PTO) -set(LLVM_TARGET_DEFINITIONS TileLibPasses.td) -mlir_tablegen(TileLibPasses.h.inc -gen-pass-decls -name PTOTileLib) - # [关键] 重命名为 PTOPassesIncGen (去掉 MLIR 前缀) add_public_tablegen_target(PTOPassesIncGen) diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index 624446d5cf..e340130281 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -133,10 +133,8 @@ std::unique_ptr createVMILegalizeArithSelectPass(); std::unique_ptr createVMILowerUnifiedToLegacyPass(); std::unique_ptr createVMINormalizeSignlessIntToUnsignedPass(); std::unique_ptr createVMIToVPTOPass(); -std::unique_ptr createInsertTemplateAttributesPass( - std::shared_ptr tileLibService); -std::unique_ptr -createExpandTileOpPass(std::shared_ptr tileLibService); +std::unique_ptr createInsertTemplateAttributesPass(); +std::unique_ptr createExpandTileOpPass(); std::unique_ptr createFoldTileBufIntrinsicsPass(); std::unique_ptr createFoldTileBufIntrinsicsPass(llvm::StringRef foldMode); std::unique_ptr createPTOCanonicalizeIRPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index ab9df2ff97..a8e2bca466 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -529,6 +529,49 @@ def PTOResolveReservedBuffers : Pass<"pto-resolve-reserved-buffers", "ModuleOp"> ]; } +def InsertTemplateAttributes + : Pass<"pto-insert-template-attributes", "ModuleOp"> { + let summary = "Attach legal PTODSL template candidates to tile operations"; + let description = [{ + Queries the process-wide PTODSL TileLib runtime for legal template + candidates and stores the compact candidate list on each tile operation as + the `candidates` attribute. Each candidate contains only id, name, + loop_depth, postupdate, and tail metadata. + }]; + let constructor = "mlir::pto::createInsertTemplateAttributesPass()"; + let dependentDialects = [ + "mlir::pto::PTODialect", + "mlir::func::FuncDialect" + ]; +} + +def ExpandTileOp : Pass<"pto-expand-tile-op", "ModuleOp"> { + let summary = "Expand tile ops into calls to TileLib template functions"; + let description = [{ + Expands tile-level operations (pto.tadd, pto.tsub, etc.) by asking the + process-wide PTODSL TileLib runtime to instantiate template libraries in + the current operation's MLIRContext. The generated template functions use + tile_buf parameters and contain vector-level implementations (pto.vecscope, + pto.vlds, pto.vadd, pto.vsts, etc.). + + Each tile op is replaced by a func.call to the generated template function, + with tile_buf operands passed directly (no type bridging). + + After this pass, the Inline pass inlines template bodies, and + FoldTileBufIntrinsics resolves tile_buf_addr / tile_valid_rows / + tile_valid_cols. + }]; + let constructor = "mlir::pto::createExpandTileOpPass()"; + let dependentDialects = [ + "mlir::pto::PTODialect", + "mlir::memref::MemRefDialect", + "mlir::arith::ArithDialect", + "mlir::func::FuncDialect", + "mlir::scf::SCFDialect", + "mlir::vector::VectorDialect" + ]; +} + def FoldTileBufIntrinsics : Pass<"pto-fold-tile-buf-intrinsics", "mlir::func::FuncOp"> { let summary = "Fold structured-view intrinsics after template inlining"; let description = [{ diff --git a/include/PTO/Transforms/TileLibPasses.td b/include/PTO/Transforms/TileLibPasses.td deleted file mode 100644 index bb2cfd095a..0000000000 --- a/include/PTO/Transforms/TileLibPasses.td +++ /dev/null @@ -1,62 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -// TileLib passes require a live compiler-session service and are intentionally -// excluded from the global textual pass registry. Their constructor strings -// suppress generated default factories; PTOAS creates them through the -// service-injecting factories declared in Passes.h. - -#ifndef MLIR_DIALECT_PTO_TILELIB_PASSES -#define MLIR_DIALECT_PTO_TILELIB_PASSES - -include "mlir/Pass/PassBase.td" - -def InsertTemplateAttributes - : Pass<"pto-insert-template-attributes", "ModuleOp"> { - let summary = "Attach legal PTODSL template candidates to tile operations"; - let description = [{ - Queries the compiler's in-process PTODSL TileLib service for legal template - candidates and stores the compact candidate list on each tile operation as - the `candidates` attribute. Each candidate contains only id, name, - loop_depth, postupdate, and tail metadata. - }]; - let constructor = "mlir::pto::createInsertTemplateAttributesPass()"; - let dependentDialects = [ - "mlir::pto::PTODialect", - "mlir::func::FuncDialect" - ]; -} - -def ExpandTileOp : Pass<"pto-expand-tile-op", "ModuleOp"> { - let summary = "Expand tile ops into calls to TileLib template functions"; - let description = [{ - Expands tile-level operations (pto.tadd, pto.tsub, etc.) by asking the - compiler's in-process PTODSL TileLib service to instantiate template - libraries in the current MLIRContext. The generated template functions use - tile_buf parameters and contain vector-level implementations (pto.vecscope, - pto.vlds, pto.vadd, pto.vsts, etc.). - - Each tile op is replaced by a func.call to the generated template function, - with tile_buf operands passed directly (no type bridging). - - After this pass, the Inline pass inlines template bodies, and - FoldTileBufIntrinsics resolves tile_buf_addr / tile_valid_rows / - tile_valid_cols. - }]; - let constructor = "mlir::pto::createExpandTileOpPass()"; - let dependentDialects = [ - "mlir::pto::PTODialect", - "mlir::memref::MemRefDialect", - "mlir::arith::ArithDialect", - "mlir::func::FuncDialect", - "mlir::scf::SCFDialect", - "mlir::vector::VectorDialect" - ]; -} - -#endif // MLIR_DIALECT_PTO_TILELIB_PASSES diff --git a/include/PTO/Transforms/TileLibService.h b/include/PTO/Transforms/TileLibService.h index 9dc4bf2859..a299cd7cd6 100644 --- a/include/PTO/Transforms/TileLibService.h +++ b/include/PTO/Transforms/TileLibService.h @@ -15,6 +15,7 @@ #include "llvm/ADT/STLFunctionalExtras.h" #include "llvm/ADT/StringRef.h" +#include #include namespace mlir::pto { @@ -51,6 +52,17 @@ class TileLibService { TileLibMaterializationCallback callback) = 0; }; +/// Process-wide access to the host TileLib implementation. The runtime owns no +/// compilation context: passes obtain the current MLIRContext from their +/// operation and pass it to TileLibService::materialize. A host binding installs +/// one service implementation for the lifetime of that runtime. +class TileLibRuntime { +public: + static void install(std::shared_ptr service); + static void uninstall(TileLibService *service); + static std::shared_ptr getService(); +}; + } // namespace mlir::pto #endif // MLIR_DIALECT_PTO_TRANSFORMS_TILELIBSERVICE_H diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index 75f2c25f85..fd238eac35 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -77,6 +77,7 @@ add_mlir_dialect_library(PTOTransforms InsertSync/InsertSyncDebug.cpp PTORematerializeFixpipeVectorQuant.cpp PTOValidateIntToPtrUses.cpp + TileLibService.cpp InsertTemplateAttributes.cpp ExpandTileOp.cpp FoldTileBufIntrinsics.cpp diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 3c8530a484..f1fbd79638 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -64,7 +64,7 @@ namespace pto { namespace func = ::mlir::func; #define GEN_PASS_DEF_EXPANDTILEOP -#include "PTO/Transforms/TileLibPasses.h.inc" +#include "PTO/Transforms/Passes.h.inc" } // namespace pto } // namespace mlir @@ -800,17 +800,7 @@ struct ExpandTileOpPass : public mlir::pto::impl::ExpandTileOpBase { using ExpandTileOpBase::ExpandTileOpBase; - explicit ExpandTileOpPass( - std::shared_ptr tileLibService) - : tileLibService(std::move(tileLibService)) {} - - ExpandTileOpPass(const ExpandTileOpPass &other) - : ExpandTileOpBase(other), - tileLibService(other.tileLibService) {} - void runOnOperation() override; - - std::shared_ptr tileLibService; }; /// Serialize a JSON array of integers. @@ -1188,8 +1178,10 @@ void ExpandTileOpPass::runOnOperation() { if (!hasExpandableOps) return; + std::shared_ptr tileLibService = + pto::TileLibRuntime::getService(); if (!tileLibService) { - mod.emitError("ExpandTileOp PTODSL backend requires an in-process service"); + mod.emitError("ExpandTileOp requires an initialized PTODSL runtime"); signalPassFailure(); return; } @@ -1210,9 +1202,8 @@ void ExpandTileOpPass::runOnOperation() { namespace mlir { namespace pto { -std::unique_ptr createExpandTileOpPass( - std::shared_ptr tileLibService) { - return std::make_unique(std::move(tileLibService)); +std::unique_ptr createExpandTileOpPass() { + return std::make_unique(); } } // namespace pto diff --git a/lib/PTO/Transforms/InsertTemplateAttributes.cpp b/lib/PTO/Transforms/InsertTemplateAttributes.cpp index 1ed474429d..98277ec0a8 100644 --- a/lib/PTO/Transforms/InsertTemplateAttributes.cpp +++ b/lib/PTO/Transforms/InsertTemplateAttributes.cpp @@ -36,7 +36,7 @@ using namespace mlir; namespace mlir { namespace pto { #define GEN_PASS_DEF_INSERTTEMPLATEATTRIBUTES -#include "PTO/Transforms/TileLibPasses.h.inc" +#include "PTO/Transforms/Passes.h.inc" } // namespace pto } // namespace mlir @@ -855,14 +855,6 @@ struct InsertTemplateAttributesPass InsertTemplateAttributesPass> { using InsertTemplateAttributesBase::InsertTemplateAttributesBase; - explicit InsertTemplateAttributesPass( - std::shared_ptr tileLibService) - : tileLibService(std::move(tileLibService)) {} - - InsertTemplateAttributesPass(const InsertTemplateAttributesPass &other) - : InsertTemplateAttributesBase(other), - tileLibService(other.tileLibService) {} - void runOnOperation() override { ModuleOp module = getOperation(); @@ -873,9 +865,11 @@ struct InsertTemplateAttributesPass }); if (tileOperations.empty()) return; + std::shared_ptr tileLibService = + pto::TileLibRuntime::getService(); if (!tileLibService) { module.emitError( - "InsertTemplateAttributes requires an in-process PTODSL service"); + "InsertTemplateAttributes requires an initialized PTODSL runtime"); return signalPassFailure(); } @@ -901,8 +895,6 @@ struct InsertTemplateAttributesPass operation->setAttr(kCandidatesAttr, *candidates); } } - - std::shared_ptr tileLibService; }; } // namespace @@ -910,10 +902,8 @@ struct InsertTemplateAttributesPass namespace mlir { namespace pto { -std::unique_ptr createInsertTemplateAttributesPass( - std::shared_ptr tileLibService) { - return std::make_unique( - std::move(tileLibService)); +std::unique_ptr createInsertTemplateAttributesPass() { + return std::make_unique(); } } // namespace pto diff --git a/lib/PTO/Transforms/TileLibService.cpp b/lib/PTO/Transforms/TileLibService.cpp new file mode 100644 index 0000000000..3efc648381 --- /dev/null +++ b/lib/PTO/Transforms/TileLibService.cpp @@ -0,0 +1,43 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include "PTO/Transforms/TileLibService.h" + +#include + +namespace { + +std::mutex &getRuntimeMutex() { + static std::mutex mutex; + return mutex; +} + +std::shared_ptr &getRuntimeService() { + static std::shared_ptr service; + return service; +} + +} // namespace + +void mlir::pto::TileLibRuntime::install( + std::shared_ptr service) { + std::lock_guard lock(getRuntimeMutex()); + getRuntimeService() = std::move(service); +} + +void mlir::pto::TileLibRuntime::uninstall(TileLibService *service) { + std::lock_guard lock(getRuntimeMutex()); + if (getRuntimeService().get() == service) + getRuntimeService().reset(); +} + +std::shared_ptr +mlir::pto::TileLibRuntime::getService() { + std::lock_guard lock(getRuntimeMutex()); + return getRuntimeService(); +} diff --git a/ptodsl/ptodsl/tilelib/_compiler_runtime.py b/ptodsl/ptodsl/tilelib/_compiler_runtime.py index bc21d1c395..3fbcaa29f2 100644 --- a/ptodsl/ptodsl/tilelib/_compiler_runtime.py +++ b/ptodsl/ptodsl/tilelib/_compiler_runtime.py @@ -65,7 +65,7 @@ def materialize( descriptor, tile_specs, context_attrs=context_attrs, - ).build_module(context=context) + ).build_module_in_context(context) module.operation.verify() return module, descriptor.name diff --git a/ptodsl/ptodsl/tilelib/_render_runtime.py b/ptodsl/ptodsl/tilelib/_render_runtime.py index 4ecbca900d..bd23dcdab0 100644 --- a/ptodsl/ptodsl/tilelib/_render_runtime.py +++ b/ptodsl/ptodsl/tilelib/_render_runtime.py @@ -176,9 +176,15 @@ def trace_entry(self, *args): rewritten(*args) # Custom golden-shaped container: single module(target_arch) + func(instance, kernel_kind). - def build_module(self, context=None): - ctx = context if context is not None else make_context() - with ctx, Location.unknown(): + def build_standalone_module(self): + """Build a module in a fresh context for standalone PTODSL use.""" + return self.build_module_in_context(make_context()) + + def build_module_in_context(self, context): + """Build a compiler-owned source module in the explicit context.""" + if context is None: + raise TypeError("compiler materialization requires an explicit context") + with context, Location.unknown(): arg_types = list(self.compute_argument_types()) module, ir_fn = self._create_instance_module(arg_types) session = self.create_session(module, ir_fn) @@ -192,6 +198,10 @@ def build_module(self, context=None): self.finalize_session(session) session.validate_final_state() self.verify_module(module) + if module.context is not context: + raise RuntimeError( + "TileLib materialization returned a module from a different context" + ) return module def _create_instance_module(self, arg_types): diff --git a/ptodsl/ptodsl/tilelib/decorator.py b/ptodsl/ptodsl/tilelib/decorator.py index 270b3a38a7..415346d56f 100644 --- a/ptodsl/ptodsl/tilelib/decorator.py +++ b/ptodsl/ptodsl/tilelib/decorator.py @@ -50,7 +50,7 @@ def __init__(self, descriptor: TileTemplate, tile_specs: dict, context_attrs=Non descriptor.name, module_factory=lambda: _TemplateTrace( descriptor, tile_specs, context_attrs=context_attrs - ).build_module(), + ).build_standalone_module(), ) self.descriptor = descriptor self.tile_specs = tile_specs diff --git a/ptodsl/tests/test_ptoas_runtime.py b/ptodsl/tests/test_ptoas_runtime.py new file mode 100644 index 0000000000..2aa7db4258 --- /dev/null +++ b/ptodsl/tests/test_ptoas_runtime.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import tempfile +import unittest +from pathlib import Path + +from ptoas import _core + + +INPUT = ( + Path(__file__).parents[2] + / "test" + / "lit" + / "vpto" + / "expand_tile_op_ptodsl_tadd.pto" +) + + +class PTOASRuntimeTest(unittest.TestCase): + def test_process_runtime_serves_consecutive_compilation_contexts(self): + self.assertTrue(INPUT.exists(), f"missing test input {INPUT}") + + with tempfile.TemporaryDirectory() as temp_dir: + for index in range(2): + output = Path(temp_dir) / f"result-{index}.mlir" + result = _core.main( + [ + "ptoas", + "--pto-arch=a5", + "--pto-backend=vpto", + "--emit-vpto", + str(INPUT), + "-o", + str(output), + ] + ) + + self.assertEqual(result, 0) + self.assertTrue(output.exists()) + self.assertIn("pto.vadd", output.read_text(encoding="utf-8")) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto b/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto index 5dd33b1c85..b7f22c355f 100644 --- a/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto +++ b/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto @@ -13,12 +13,11 @@ // Running without --enable-op-fusion proves metadata insertion is not gated // by fusion. Printing before FusionPlan proves its position when fusion is on. // -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after-all %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --pto-level=level2 --enable-op-fusion --emit-pto-ir --mlir-print-ir-before=pto-fusion-plan %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=PREFUSION -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after-all %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>/dev/null | FileCheck %s --check-prefix=EXPAND -// META: IR Dump After InsertTemplateAttributes // META: pto.tadd // META-SAME: candidates = [ // META-SAME: id = 0 : i64 @@ -34,7 +33,6 @@ // PREFUSION: pto.tadd // PREFUSION-SAME: candidates = [ -// SELECT: IR Dump After ExpandTileOp // SELECT: func.func {{.*}}@{{.*}}__template_tadd // EXPAND: func.func @TADD diff --git a/test/lit/vpto/vpto_pipeline_vmi_after_tileop_expand.pto b/test/lit/vpto/vpto_pipeline_vmi_after_tileop_expand.pto index 92ecb930b1..a87c473260 100644 --- a/test/lit/vpto/vpto_pipeline_vmi_after_tileop_expand.pto +++ b/test/lit/vpto/vpto_pipeline_vmi_after_tileop_expand.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o /dev/null --mlir-print-ir-after-all 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o /dev/null --mlir-print-ir-after=pto-expand-tile-op --mlir-print-ir-after=vmi-to-vpto --mlir-print-ir-after=vpto-ptr-normalize --mlir-print-ir-after=pto-infer-vpto-vecscope --mlir-print-ir-after=loop-invariant-code-motion 2>&1 | FileCheck %s // CHECK: IR Dump After ExpandTileOp // CHECK: IR Dump After VMIToVPTO diff --git a/tools/ptoas/NativeModule.cpp b/tools/ptoas/NativeModule.cpp index ff81286144..0f0bda26d4 100644 --- a/tools/ptoas/NativeModule.cpp +++ b/tools/ptoas/NativeModule.cpp @@ -27,14 +27,11 @@ namespace { class PythonTileLibService final : public mlir::pto::TileLibService { public: - explicit PythonTileLibService(py::object contextOwner) - : contextOwner(std::move(contextOwner)) {} - mlir::FailureOr getMetadata(const mlir::pto::TileLibMaterializationRequest &request) override { py::gil_scoped_acquire acquire; try { - return py::cast(getRuntime().attr("metadata")( + return py::cast(getCompilerRuntime().attr("metadata")( request.target, request.op, request.operandSpecsJson, request.contextAttrsJson)); } catch (const py::error_already_set &error) { @@ -51,6 +48,7 @@ class PythonTileLibService final : public mlir::pto::TileLibService { mlir::pto::TileLibMaterializationCallback callback) override { py::gil_scoped_acquire acquire; try { + py::object contextOwner = getPythonContext(context); MlirContext pythonContext = py::cast(contextOwner); if (unwrap(pythonContext) != &context) { llvm::errs() << "TileLib: Python context does not match the PTOAS " @@ -58,7 +56,7 @@ class PythonTileLibService final : public mlir::pto::TileLibService { return mlir::failure(); } - py::tuple result = getRuntime().attr("materialize")( + py::tuple result = getCompilerRuntime().attr("materialize")( request.target, request.op, request.operandSpecsJson, request.contextAttrsJson, request.candidateId, contextOwner); if (result.size() != 2) @@ -90,19 +88,49 @@ class PythonTileLibService final : public mlir::pto::TileLibService { } private: - py::object &getRuntime() { - if (!runtime) - runtime = py::module_::import("ptodsl.tilelib._compiler_runtime"); - return runtime; + static py::module_ getCompilerRuntime() { + // Python's sys.modules cache makes this a process-wide runtime module + // without storing a py::object whose destructor could outlive CPython. + return py::module_::import("ptodsl.tilelib._compiler_runtime"); + } + + static py::object getPythonContext(mlir::MLIRContext &context) { + py::object capsule = py::reinterpret_steal( + mlirPythonContextToCapsule(wrap(&context))); + return py::module_::import("ptoas.mlir.ir") + .attr("Context") + .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule); + } +}; + +constexpr char kRuntimeRegistrationCapsuleName[] = + "ptoas.TileLibRuntimeRegistration"; + +class PythonTileLibRuntimeRegistration { +public: + PythonTileLibRuntimeRegistration() + : service(std::make_shared()) { + mlir::pto::TileLibRuntime::install(service); + } + + ~PythonTileLibRuntimeRegistration() { + mlir::pto::TileLibRuntime::uninstall(service.get()); } - // These objects are created and destroyed by runPTOASFromPython while the - // calling thread owns the GIL. materialize() reacquires it for every DSL - // invocation because the native compiler releases it around the driver. - py::object contextOwner; - py::object runtime; +private: + std::shared_ptr service; }; +void destroyRuntimeRegistration(PyObject *capsule) { + void *pointer = + PyCapsule_GetPointer(capsule, kRuntimeRegistrationCapsuleName); + if (!pointer) { + PyErr_Clear(); + return; + } + delete static_cast(pointer); +} + int runPTOASFromPython(const std::vector &arguments) { std::vector storage = arguments; std::vector argv; @@ -113,14 +141,12 @@ int runPTOASFromPython(const std::vector &arguments) { py::object contextOwner = py::module_::import("ptoas.mlir.ir").attr("Context")(); MlirContext rawContext = py::cast(contextOwner); - auto tileLibService = - std::make_shared(contextOwner); int result; { py::gil_scoped_release release; result = mlir::pto::runPTOAS(static_cast(argv.size()), argv.data(), - *unwrap(rawContext), tileLibService); + *unwrap(rawContext)); } return result; } @@ -131,5 +157,10 @@ PYBIND11_MODULE(_core, module) { module.doc() = "PTOAS compiler and PTO dialect native bindings"; py::module_::import("ptoas.mlir.ir"); mlir::pto::python::populatePTODialectBindings(module); + module.add_object( + "_tilelib_runtime_registration", + py::capsule(new PythonTileLibRuntimeRegistration(), + kRuntimeRegistrationCapsuleName, + destroyRuntimeRegistration)); module.def("main", &runPTOASFromPython, py::arg("argv")); } diff --git a/tools/ptoas/driver.cpp b/tools/ptoas/driver.cpp index 5f4320b992..d9b06c4249 100644 --- a/tools/ptoas/driver.cpp +++ b/tools/ptoas/driver.cpp @@ -683,11 +683,10 @@ mlir::pto::PTOASContext::PTOASContext(DialectRegistry ®istry, argc(argc), argv(argv) {} mlir::pto::PTOASContext::PTOASContext( - MLIRContext &borrowedContext, - std::shared_ptr tileLibService, - llvm::StringRef outputPath, int argc, char **argv) - : mlirContext(&borrowedContext), tileLibService(std::move(tileLibService)), - outputPath(outputPath.str()), argc(argc), argv(argv) {} + MLIRContext &borrowedContext, llvm::StringRef outputPath, int argc, + char **argv) + : mlirContext(&borrowedContext), outputPath(outputPath.str()), argc(argc), + argv(argv) {} mlir::pto::PTOASContext::~PTOASContext() = default; @@ -708,11 +707,6 @@ void mlir::pto::PTOASContext::initializeMLIRContext() { MLIRContext &mlir::pto::PTOASContext::getMLIRContext() { return *mlirContext; } -std::shared_ptr -mlir::pto::PTOASContext::getTileLibService() const { - return tileLibService; -} - void mlir::pto::PTOASContext::setArch(std::string value) { arch = std::move(value); } @@ -1279,9 +1273,8 @@ static LogicalResult writeTextOutput(llvm::StringRef output, // +-------------+ +------------------------------------------+ // | C++ source | | fatobj | // +-------------+ +------------------------------------------+ -static int runPTOASDriver( - int argc, char **argv, MLIRContext *borrowedContext = nullptr, - std::shared_ptr tileLibService = nullptr) { +static int runPTOASDriver(int argc, char **argv, + MLIRContext *borrowedContext = nullptr) { DialectRegistry registry; mlir::pto::registerPTOASDialects(registry); if (borrowedContext) @@ -1302,8 +1295,8 @@ static int runPTOASDriver( std::unique_ptr context; if (borrowedContext) { - context = std::make_unique( - *borrowedContext, std::move(tileLibService), outputFilename, argc, argv); + context = std::make_unique(*borrowedContext, outputFilename, + argc, argv); } else { context = std::make_unique(registry, outputFilename, argc, argv); @@ -1347,9 +1340,7 @@ int mlir::pto::runPTOAS(int argc, char **argv) { return runPTOASDriver(argc, argv); } -int mlir::pto::runPTOAS( - int argc, char **argv, MLIRContext &borrowedContext, - std::shared_ptr tileLibService) { - return runPTOASDriver(argc, argv, &borrowedContext, - std::move(tileLibService)); +int mlir::pto::runPTOAS(int argc, char **argv, + MLIRContext &borrowedContext) { + return runPTOASDriver(argc, argv, &borrowedContext); } diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index fd00bbce2c..a6b59bd761 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -2787,9 +2787,7 @@ static void prepareVPTOForEmission(PassManager &pm) { kernelModulePM.addPass(pto::createPTOValidateVPTOEmissionIRPass()); } -static void -lowerPTOToVPTOBackend(PassManager &pm, ModuleOp module, - std::shared_ptr tileLibService) { +static void lowerPTOToVPTOBackend(PassManager &pm, ModuleOp module) { auto &kernelModulePM = pm.nest(); auto moduleArchAttr = module->getAttrOfType("pto.target_arch"); @@ -2806,8 +2804,7 @@ lowerPTOToVPTOBackend(PassManager &pm, ModuleOp module, return; } - kernelModulePM.addPass( - pto::createExpandTileOpPass(std::move(tileLibService))); + kernelModulePM.addPass(pto::createExpandTileOpPass()); kernelModulePM.addPass(pto::createPTOInlineLibCallPass()); kernelModulePM.addNestedPass( @@ -2902,15 +2899,13 @@ static int emitVPTOBackendResult(ModuleOp module, PTOASCompileResult &result, } static LogicalResult runVPTOBackendPipeline(OwningOpRef &module, - bool hasTileOpsToExpand, - std::shared_ptr - tileLibService) { + bool hasTileOpsToExpand) { PassManager pm(module->getContext()); pm.enableVerifier(); pm.addPass(pto::createVPTOSplitCVModulePass()); pm.addPass(pto::createVPTONormalizeContainerPass()); if (hasTileOpsToExpand) - lowerPTOToVPTOBackend(pm, module.get(), std::move(tileLibService)); + lowerPTOToVPTOBackend(pm, module.get()); auto &kernelModulePM = pm.nest(); // Inline legal direct calls before VMI layout assignment so private helper // bodies participate in one caller-local layout decision. The Func @@ -3211,8 +3206,7 @@ int mlir::pto::compilePTOASModule( "skipping the shared PTO-to-VPTO lowering pipeline.\n"; return 1; } - if (failed(runVPTOBackendPipeline(module, hasTileOpsToExpand, - context.getTileLibService()))) + if (failed(runVPTOBackendPipeline(module, hasTileOpsToExpand))) return 1; return emitVPTOBackendResult(*module, result, emitVPTOHostStub, context.getCANNVersionOrDefault()); @@ -3251,8 +3245,7 @@ int mlir::pto::compilePTOASModule( // Fusion may later filter the ordered `candidates` array; ExpandTileOp // consumes the first candidate that remains. if (!isA2A3 && effectiveBackend == PTOBackend::VPTO && hasTileOpsToExpand) - pm.addPass(pto::createInsertTemplateAttributesPass( - context.getTileLibService())); + pm.addPass(pto::createInsertTemplateAttributesPass()); // Keep frontend fusion on tile-native PTO IR and annotate last_use directly // on scheduled block-local spans before the shared mainline lowers tiles. @@ -3395,8 +3388,7 @@ int mlir::pto::compilePTOASModule( if (failed(emitSharedPreBackendSeamIR(*module, ptoSeamIRFile))) return 1; - if (failed(runVPTOBackendPipeline( - module, hasTileOpsToExpand, context.getTileLibService()))) + if (failed(runVPTOBackendPipeline(module, hasTileOpsToExpand))) return 1; return emitVPTOBackendResult(*module, result, emitVPTOHostStub, context.getCANNVersionOrDefault()); diff --git a/tools/ptoas/ptoas.h b/tools/ptoas/ptoas.h index b2d317b043..500525094c 100644 --- a/tools/ptoas/ptoas.h +++ b/tools/ptoas/ptoas.h @@ -13,7 +13,6 @@ #include "PTO/Compiler/CompilerApi.h" #include "PTO/Transforms/VPTOLLVMEmitter.h" #include "VFSIMTSizePatcher.h" -#include "PTO/Transforms/TileLibService.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/MLIRContext.h" #include "mlir/Support/LogicalResult.h" @@ -63,9 +62,8 @@ class PTOASContext { public: PTOASContext(DialectRegistry ®istry, llvm::StringRef outputPath, int argc, char **argv); - PTOASContext(MLIRContext &borrowedContext, - std::shared_ptr tileLibService, - llvm::StringRef outputPath, int argc, char **argv); + PTOASContext(MLIRContext &borrowedContext, llvm::StringRef outputPath, + int argc, char **argv); ~PTOASContext(); LogicalResult initializeEnvironment(bool requiresToolchain, @@ -73,7 +71,6 @@ class PTOASContext { void initializeMLIRContext(); MLIRContext &getMLIRContext(); - std::shared_ptr getTileLibService() const; void setArch(std::string value); llvm::StringRef getArch() const; @@ -101,7 +98,6 @@ class PTOASContext { private: std::unique_ptr ownedMlirContext; MLIRContext *mlirContext = nullptr; - std::shared_ptr tileLibService; std::string outputPath; std::string arch; BackendInfo backendInfo; @@ -143,8 +139,7 @@ void loadPTOASDialects(MLIRContext &context); // Reusable driver entry shared by the Python extension and standalone CLI. PTOAS_COMPILER_EXPORT int runPTOAS(int argc, char **argv); PTOAS_COMPILER_EXPORT int -runPTOAS(int argc, char **argv, MLIRContext &borrowedContext, - std::shared_ptr tileLibService); +runPTOAS(int argc, char **argv, MLIRContext &borrowedContext); // Attach textual-.pto SSA name hints (function args, block args, op results) // to the parsed module's Locations as debug metadata. Called by the driver From 988d50e245217669a27448c96641bb7eaf26baed Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Wed, 5 Aug 2026 15:00:15 +0800 Subject: [PATCH 072/122] fix(ptoas): reuse embedded compiler state safely --- lib/PTO/Transforms/ExpandTileOp.cpp | 31 +++++++----- ptodsl/tests/test_ptoas_runtime.py | 53 ++++++++++++++++++-- test/lit/vpto/expand_tile_op_ptodsl_tadd.pto | 20 ++++++-- tools/ptoas/driver.cpp | 5 ++ 4 files changed, 90 insertions(+), 19 deletions(-) diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index f1fbd79638..8a255473b7 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -786,8 +786,9 @@ struct ExpandState { func::FuncOp invokeTileLib(const SpecKey &key, Operation *tileOp, ModuleOp mod, MLIRContext *ctx); func::FuncOp invokeInProcessTileLib(const SpecKey &key, - StringRef candidateId, ModuleOp mod, - MLIRContext *ctx); + StringRef candidateId, + const std::string &uniqueName, + ModuleOp mod, MLIRContext *ctx); LogicalResult expandTileOpsInFunction(func::FuncOp func, ModuleOp mod, MLIRContext *ctx); @@ -947,6 +948,14 @@ static std::string buildUniqueFunctionBaseName(const SpecKey &key) { return uniqueName; } +static std::string buildUniqueFunctionName(const SpecKey &key, + StringRef candidateId) { + std::string uniqueName = buildUniqueFunctionBaseName(key); + if (!candidateId.empty()) + uniqueName += "__" + candidateId.str(); + return uniqueName; +} + static std::string buildContextAttrsJson(const SpecKey &key) { std::string json = "{"; for (size_t i = 0; i < key.contextAttrs.size(); ++i) { @@ -970,6 +979,7 @@ static std::string buildContextAttrsJson(const SpecKey &key) { // ============================================================================ func::FuncOp ExpandState::invokeInProcessTileLib(const SpecKey &key, StringRef candidateId, + const std::string &uniqueName, ModuleOp mod, MLIRContext *ctx) { if (!tileLibService) @@ -1006,16 +1016,7 @@ func::FuncOp ExpandState::invokeInProcessTileLib(const SpecKey &key, return failure(); } - std::string uniqueName = buildUniqueFunctionBaseName(key); - if (!candidateId.empty()) - uniqueName += "__" + candidateId.str(); - SymbolTable targetSymTable(mod); - if (auto existingFunc = targetSymTable.lookup(uniqueName)) { - importedEntry = cast(existingFunc); - return success(); - } - llvm::StringMap plannedSymbols; for (func::FuncOp fn : sourceFuncs) { std::string newName = fn == sourceEntry @@ -1103,7 +1104,13 @@ func::FuncOp ExpandState::invokeTileLib(const SpecKey &key, return nullptr; } - return invokeInProcessTileLib(key, selectedName.getValue(), mod, ctx); + std::string uniqueName = + buildUniqueFunctionName(key, selectedName.getValue()); + if (auto existing = mod.lookupSymbol(uniqueName)) + return existing; + + return invokeInProcessTileLib(key, selectedName.getValue(), uniqueName, mod, + ctx); } // ============================================================================ diff --git a/ptodsl/tests/test_ptoas_runtime.py b/ptodsl/tests/test_ptoas_runtime.py index 2aa7db4258..d8a03c23d7 100644 --- a/ptodsl/tests/test_ptoas_runtime.py +++ b/ptodsl/tests/test_ptoas_runtime.py @@ -10,8 +10,10 @@ import tempfile import unittest from pathlib import Path +from unittest import mock from ptoas import _core +from ptodsl.tilelib import _compiler_runtime INPUT = ( @@ -28,14 +30,19 @@ def test_process_runtime_serves_consecutive_compilation_contexts(self): self.assertTrue(INPUT.exists(), f"missing test input {INPUT}") with tempfile.TemporaryDirectory() as temp_dir: - for index in range(2): - output = Path(temp_dir) / f"result-{index}.mlir" + pto_output = Path(temp_dir) / "result-pto.mlir" + vpto_output = Path(temp_dir) / "result-vpto.mlir" + + for output_mode, output in ( + ("--emit-pto-ir", pto_output), + ("--emit-vpto", vpto_output), + ): result = _core.main( [ "ptoas", "--pto-arch=a5", "--pto-backend=vpto", - "--emit-vpto", + output_mode, str(INPUT), "-o", str(output), @@ -44,7 +51,45 @@ def test_process_runtime_serves_consecutive_compilation_contexts(self): self.assertEqual(result, 0) self.assertTrue(output.exists()) - self.assertIn("pto.vadd", output.read_text(encoding="utf-8")) + + pto_ir = pto_output.read_text(encoding="utf-8") + vpto_ir = vpto_output.read_text(encoding="utf-8") + self.assertIn("pto.tadd ins", pto_ir) + self.assertNotIn("pto.vadd", pto_ir) + self.assertNotIn("pto.tadd ins", vpto_ir) + self.assertIn("pto.vadd", vpto_ir) + + def test_reuses_imported_specialization_before_materializing_again(self): + calls = 0 + original_materialize = _compiler_runtime.materialize + + def counted_materialize(*args, **kwargs): + nonlocal calls + calls += 1 + return original_materialize(*args, **kwargs) + + with tempfile.TemporaryDirectory() as temp_dir: + output = Path(temp_dir) / "result-vpto.mlir" + with mock.patch.object( + _compiler_runtime, + "materialize", + side_effect=counted_materialize, + ): + result = _core.main( + [ + "ptoas", + "--pto-arch=a5", + "--pto-backend=vpto", + "--emit-vpto", + str(INPUT), + "-o", + str(output), + ] + ) + + self.assertEqual(result, 0) + self.assertEqual(calls, 1) + self.assertIn("pto.vadd", output.read_text(encoding="utf-8")) if __name__ == "__main__": diff --git a/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto b/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto index b7f22c355f..fbef117298 100644 --- a/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto +++ b/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto @@ -33,7 +33,8 @@ // PREFUSION: pto.tadd // PREFUSION-SAME: candidates = [ -// SELECT: func.func {{.*}}@{{.*}}__template_tadd +// SELECT-COUNT-2: call @{{.*}}__template_tadd +// SELECT-COUNT-1: func.func {{.*}}@{{.*}}__template_tadd // EXPAND: func.func @TADD // EXPAND-NOT: pto.tadd ins @@ -52,7 +53,10 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { %b = pto.alloc_tile : !pto.tile_buf - %dst = pto.alloc_tile + %dst0 = pto.alloc_tile + : !pto.tile_buf + %dst1 = pto.alloc_tile : !pto.tile_buf @@ -63,7 +67,17 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { !pto.tile_buf) outs( - %dst + %dst0 + : !pto.tile_buf) + pto.tadd ins( + %a, %b + : !pto.tile_buf, + !pto.tile_buf) + outs( + %dst1 : !pto.tile_buf) return diff --git a/tools/ptoas/driver.cpp b/tools/ptoas/driver.cpp index d9b06c4249..06644f4307 100644 --- a/tools/ptoas/driver.cpp +++ b/tools/ptoas/driver.cpp @@ -1282,6 +1282,11 @@ static int runPTOASDriver(int argc, char **argv, mlir::pto::registerPTOASPassesAndCLOptions(); llvm::cl::SetVersionPrinter(printPTOASVersion); + // The Python entry point may invoke the driver repeatedly in one process. + // Restore every registered LLVM option to its declared default before + // parsing the next invocation. + llvm::cl::ResetAllOptionOccurrences(); + const bool cliArchSpecified = hasCLIOption(argc, argv, "--pto-arch"); const bool cliBackendSpecified = hasCLIOption(argc, argv, "--pto-backend"); From 0c6d8b2e74764f27d064c3c31344d38ca6da9214 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 7 Aug 2026 17:19:43 +0800 Subject: [PATCH 073/122] refactor(emitc): fix #1165 at the source via peel in MGATHER/MSCATTER Per PR review: prefer fixing at the source over the sink. MGATHER/MSCATTER are template intrinsics that accept the concrete descriptor directly, so peel the type-converter materialization bridge on their mem/idx/dst (src) operands. The static-stride GlobalTensor bridge then becomes dead and is dropped by the first cast-cleanup rule (use_empty), so it never reaches the emitc.cast fallback and no invalid C-style GlobalTensor<...> cast is emitted. This removes the sink-side areRefinableGlobalTensorTypes helper and its forward/feedsReturn branch, which only tolerated the inconsistency at the exit. Bringing MGATHER/MSCATTER in line with the peel convention already used by the other operand-consuming patterns keeps the consistency contract in one place. --- lib/PTO/Transforms/PTOToEmitC.cpp | 95 +++++-------------------------- 1 file changed, 14 insertions(+), 81 deletions(-) diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index d26ca6a7f4..c54a52eb66 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -3412,9 +3412,14 @@ struct PTOMGatherToMGATHER : public OpConversionPattern { LogicalResult matchAndRewrite(pto::MGatherOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { auto *ctx = rewriter.getContext(); - Value mem = adaptor.getMem(); - Value idx = adaptor.getIdx(); - Value dst = adaptor.getDst(); + // MGATHER is a template intrinsic that accepts the concrete descriptor + // directly, so peel any type-converter materialization bridge and feed the + // producing value. This keeps the compile-time static-stride GlobalTensor + // from the partition_view pattern instead of the dynamic-stride bridge, + // whose GlobalTensor<...> C-style cast would not compile (issue #1165). + Value mem = peelUnrealized(adaptor.getMem()); + Value idx = peelUnrealized(adaptor.getIdx()); + Value dst = peelUnrealized(adaptor.getDst()); Value memArg = mem; auto coalescePropAttr = @@ -6169,9 +6174,12 @@ struct PTOMScatterToMSCATTER : public OpConversionPattern { LogicalResult matchAndRewrite(pto::MScatterOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { auto *ctx = rewriter.getContext(); - Value src = adaptor.getSrc(); - Value idx = adaptor.getIdx(); - Value mem = adaptor.getMem(); + // MSCATTER is a template intrinsic that accepts the concrete descriptor + // directly, so peel any type-converter materialization bridge and feed the + // producing value (static-stride GlobalTensor). See MGATHER above / #1165. + Value src = peelUnrealized(adaptor.getSrc()); + Value idx = peelUnrealized(adaptor.getIdx()); + Value mem = peelUnrealized(adaptor.getMem()); auto coalesceAttr = dyn_cast_or_null(op.getProperties().coalesce); auto scatterAtomicAttr = @@ -13669,61 +13677,6 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, populateBranchOpInterfaceTypeConversionPattern(patterns, typeConverter); } -// A cast between two `GlobalTensor<...>` opaque C++ types that are identical -// except that one side carries concrete Shape/Stride template values while the -// other uses the fully-dynamic `-1` placeholders (as produced by the -// PTOToEmitCTypeConverter, which cannot recover strides from the stride-less -// tensor_view type) has no valid C++ converting constructor. The values are -// interchangeable at every templated backend call site (e.g. MGATHER/MSCATTER), -// so such a bridge must forward the value rather than lower to an invalid -// C-style `emitc.cast`. -static bool areRefinableGlobalTensorTypes(Type a, Type b) { - auto oa = dyn_cast(a); - auto ob = dyn_cast(b); - if (!oa || !ob) - return false; - StringRef sa = oa.getValue(); - StringRef sb = ob.getValue(); - if (!sa.contains("GlobalTensor<") || !sb.contains("GlobalTensor<")) - return false; - - SmallVector shapeA, shapeB, strideA, strideB; - if (!parseIntegerTemplateList(sa, "Shape<", shapeA) || - !parseIntegerTemplateList(sb, "Shape<", shapeB) || - !parseIntegerTemplateList(sa, "Stride<", strideA) || - !parseIntegerTemplateList(sb, "Stride<", strideB)) - return false; - - // Element-wise compatible if equal or one side is the `-1` wildcard. - auto listsRefinable = [](ArrayRef x, ArrayRef y) { - if (x.size() != y.size()) - return false; - for (auto [u, v] : llvm::zip(x, y)) - if (u != v && u != -1 && v != -1) - return false; - return true; - }; - if (!listsRefinable(shapeA, shapeB) || !listsRefinable(strideA, strideB)) - return false; - - // Everything outside the Shape<...>/Stride<...> lists (element type, layout, - // overall structure) must match exactly. - auto blank = [](StringRef s, StringRef marker) -> std::string { - std::string out = s.str(); - size_t pos = out.find(marker.str()); - if (pos == std::string::npos) - return out; - size_t start = pos + marker.size(); - size_t end = out.find('>', start); - if (end == std::string::npos) - return out; - out.erase(start, end - start); - return out; - }; - return blank(blank(sa, "Shape<"), "Stride<") == - blank(blank(sb, "Shape<"), "Stride<"); -} - //===----------------------------------------------------------------------===// // Pass //===----------------------------------------------------------------------===// @@ -14146,26 +14099,6 @@ static AICORE inline void PTOAS__DCCI_SINGLE_CACHE_LINE(Ptr ptr) { return; } - // A static-stride `GlobalTensor` produced by the partition_view static - // pattern and the fully-dynamic-stride `GlobalTensor` demanded by the - // type converter have no C++ converting constructor. Forwarding the - // refined (static) value is only safe when every consumer accepts a - // more-specific GlobalTensor template instantiation (e.g. MGATHER / - // MSCATTER, which are C++ templates). A `return` must match the enclosing - // function's fixed (dynamic) result type exactly, so forwarding there - // would break verification -- fall through to the emitc.cast branch in - // that case. - if (areRefinableGlobalTensorTypes(inTy, outTy)) { - bool feedsReturn = llvm::any_of(output.getUsers(), [](Operation *user) { - return isa(user); - }); - if (!feedsReturn) { - output.replaceAllUsesWith(input); - castsToErase.push_back(cast); - return; - } - } - if (emitc::isSupportedEmitCType(inTy) && emitc::isSupportedEmitCType(outTy)) { OpBuilder builder(cast); auto c = builder.create(cast.getLoc(), outTy, input); From a1aa09e510fa744fe48657a80fd2c817d1b74951 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 8 Aug 2026 11:01:26 +0800 Subject: [PATCH 074/122] Align A5 ttrans implicit tmp with other ops TTrans passed a bare requireExplicitTmp, causing A5+level3 to wrongly error on a missing tmp. The A5 backend ignores the ttrans tmp buffer, so materialize a no-address placeholder like the other A5 ops instead. --- .../Transforms/PTOMaterializeImplicitTmp.cpp | 11 ++++--- .../lit/pto/ttrans_implicit_tmp_a5_level3.pto | 30 +++++++++++++++++++ 2 files changed, 37 insertions(+), 4 deletions(-) create mode 100644 test/lit/pto/ttrans_implicit_tmp_a5_level3.pto diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp index b832503501..92139a09e6 100644 --- a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -627,6 +627,8 @@ static LogicalResult materializeFixedMandatoryTmp(Operation *op, .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) return success(); + bool isA5 = + pto::getTargetArch(op) == pto::PTOArch::A5; auto srcTy = dyn_cast(typedOp.getSrc().getType()); auto dstTy = dyn_cast(typedOp.getDst().getType()); auto srcShape = getShapeVec(typedOp.getSrc().getType()); @@ -643,16 +645,17 @@ static LogicalResult materializeFixedMandatoryTmp(Operation *op, bool usesTmp = dstShape[1] % rowStride == 0 && srcShape[1] % elemPerBlock == 0 && srcShape[1] / elemPerBlock <= 255; - FailureOr type = makeSameShapeTmpType( - ctx, typedOp.getSrc()); - if (!usesTmp) + FailureOr type = + isA5 ? makeA5PlaceholderTmpType(ctx, typedOp.getSrc()) + : makeSameShapeTmpType(ctx, typedOp.getSrc()); + if (!isA5 && !usesTmp) type = makeVecTmpType(ctx, {1, elemPerBlock}, srcTy.getElementType(), {1, elemPerBlock}); if (failed(type)) return typedOp.emitOpError("failed to build implicit ttrans tmp"); return replaceFixedDpsOpWithTmp( op, {typedOp.getSrc(), Value(), typedOp.getDst()}, *type, - {1, 1, 1}, requireExplicitTmp, "ttrans"); + {1, 1, 1}, isA5 ? false : requireExplicitTmp, "ttrans"); }) .Default([](Operation *) { return success(); }); } diff --git a/test/lit/pto/ttrans_implicit_tmp_a5_level3.pto b/test/lit/pto/ttrans_implicit_tmp_a5_level3.pto new file mode 100644 index 0000000000..98ccbe2945 --- /dev/null +++ b/test/lit/pto/ttrans_implicit_tmp_a5_level3.pto @@ -0,0 +1,30 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// On A5 the backend ignores the ttrans tmp buffer, so an implicit tmp must be +// materialized as a no-address placeholder at both level2 and level3 rather +// than rejected. level3 skips PlanMemory; the ttrans handler must not require +// an explicit tmp on A5 (consistent with the other A5 ops). + +// RUN: sed -E 's/ addr = %[A-Za-z0-9_]+//g' %s > %t.level2.pto && ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %t.level2.pto 2>&1 | FileCheck %s --check-prefix=A5 +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 + +module { + func.func @a5_ttrans_implicit_tmp() attributes {pto.kernel_kind = #pto.kernel_kind} { + %a0 = arith.constant 0 : i64 + %a2048 = arith.constant 2048 : i64 + %src = pto.alloc_tile addr = %a0 : !pto.tile_buf + %dst = pto.alloc_tile addr = %a2048 : !pto.tile_buf + pto.ttrans ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// A5-LABEL: func.func @a5_ttrans_implicit_tmp +// A5: pto.ttrans ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) From e2014da8fd86f786bd48f9b50d4b0224d5f7c315 Mon Sep 17 00:00:00 2001 From: and0d0 Date: Sun, 9 Aug 2026 21:30:43 +0800 Subject: [PATCH 075/122] docs: document explicit L1-to-L0 load limits --- ptodsl/docs/user_guide/07-data-movement-ops.md | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/ptodsl/docs/user_guide/07-data-movement-ops.md b/ptodsl/docs/user_guide/07-data-movement-ops.md index 00ba51b3e1..7fffb58a80 100644 --- a/ptodsl/docs/user_guide/07-data-movement-ops.md +++ b/ptodsl/docs/user_guide/07-data-movement-ops.md @@ -944,17 +944,21 @@ Cube compute step; it does not issue those transfers itself. **Description**: Explicit-control L1-to-L0A/L0B loads. This overload preserves the authored fractal-block control fields and does not infer strides from a -logical tile shape. +logical tile shape. It currently rejects FP4 packed source pointers because +the compatibility wrapper cannot yet guarantee selection of the FP4-specific +L1-to-L0 intrinsic on this path. For FP4 packed operands, use the structured +shape-derived `mte_l1_l0a(..., m, k, ...)` or `mte_l1_l0b(..., k, n, ...)` +form below. **Parameters**: | Parameter | Type | Description | |-----------|------|-------------| -| `src` | `PtrType` (L1/MAT) | L1 source pointer; parent-allocation matrix offsets are represented by the control fields below. | +| `src` | `PtrType` (L1/MAT) | L1 source pointer; parent-allocation matrix offsets are represented by the control fields below. FP4 packed source pointers are not supported by this explicit-control overload. | | `dst` | `PtrType` (L0A/L0B) | L0 destination pointer; stage/version offsets belong in this pointer. | -| `m_start`, `k_start` | `int` | Source fractal-block coordinates at which the load begins. | -| `m_step`, `k_step` | `int` | Number of fractal blocks transferred along the two source axes. | -| `src_stride`, `dst_stride` | `int` | Physical outer strides of the complete L1 and L0 allocations, in fractal-block units. They are independent of the transferred region extents. | +| `m_start`, `k_start` | `int` in `0..65535` | Source fractal-block coordinates at which the load begins. | +| `m_step`, `k_step` | `int` in `1..255` | Number of fractal blocks transferred along the two source axes. | +| `src_stride`, `dst_stride` | `int` in `1..65535` | Physical outer strides of the complete L1 and L0 allocations, in fractal-block units. They are independent of the transferred region extents. | | `transpose` | `bool` | Final hardware transpose attribute. | **Returns**: None (side-effect operation). From 3d51925debd46004fbb3d9bbb9b532c854639afd Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Mon, 10 Aug 2026 09:36:34 +0800 Subject: [PATCH 076/122] fix: make TileOps editable --- ptodsl/tests/test_python_package_layout.py | 27 ++++++++++++++++++++++ pyproject.toml | 1 + 2 files changed, 28 insertions(+) create mode 100644 ptodsl/tests/test_python_package_layout.py diff --git a/ptodsl/tests/test_python_package_layout.py b/ptodsl/tests/test_python_package_layout.py new file mode 100644 index 0000000000..a8ac851fc4 --- /dev/null +++ b/ptodsl/tests/test_python_package_layout.py @@ -0,0 +1,27 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software; you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import tomllib +import unittest +from pathlib import Path + + +class PythonPackageLayoutTest(unittest.TestCase): + def test_tileops_is_declared_as_an_editable_python_package(self): + project_root = Path(__file__).resolve().parents[2] + with (project_root / "pyproject.toml").open("rb") as stream: + pyproject = tomllib.load(stream) + + packages = pyproject["tool"]["scikit-build"]["wheel"]["packages"] + self.assertEqual(packages["TileOps"], "lib/TileOps") + self.assertTrue((project_root / packages["TileOps"] / "__init__.py").is_file()) + + +if __name__ == "__main__": + unittest.main() diff --git a/pyproject.toml b/pyproject.toml index 05434dce5e..56ecc54767 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ install.components = ["PTOAS_Python"] [tool.scikit-build.wheel.packages] ptodsl = "ptodsl/ptodsl" ptoas = "ptodsl/ptoas" +TileOps = "lib/TileOps" [tool.scikit-build.cmake.define] PTOAS_RELEASE_VERSION_OVERRIDE = { env = "PTOAS_RELEASE_VERSION_OVERRIDE", default = "" } From dfa0a8cd33cbf5b3a54a8c7b35902e6e99b3fe25 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Mon, 10 Aug 2026 09:52:11 +0800 Subject: [PATCH 077/122] test: support Python 3.10 package check --- ptodsl/tests/test_python_package_layout.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/ptodsl/tests/test_python_package_layout.py b/ptodsl/tests/test_python_package_layout.py index a8ac851fc4..4e737ed471 100644 --- a/ptodsl/tests/test_python_package_layout.py +++ b/ptodsl/tests/test_python_package_layout.py @@ -1,13 +1,11 @@ -#!/usr/bin/env python3 # Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software; you can redistribute it and/or modify it under the terms and conditions of +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of # CANN Open Software License Agreement Version 2.0 (the "License"). # Please refer to the License for details. You may not use this file except in compliance with the License. # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -import tomllib import unittest from pathlib import Path @@ -15,12 +13,11 @@ class PythonPackageLayoutTest(unittest.TestCase): def test_tileops_is_declared_as_an_editable_python_package(self): project_root = Path(__file__).resolve().parents[2] - with (project_root / "pyproject.toml").open("rb") as stream: - pyproject = tomllib.load(stream) + pyproject = (project_root / "pyproject.toml").read_text(encoding="utf-8") - packages = pyproject["tool"]["scikit-build"]["wheel"]["packages"] - self.assertEqual(packages["TileOps"], "lib/TileOps") - self.assertTrue((project_root / packages["TileOps"] / "__init__.py").is_file()) + self.assertIn("[tool.scikit-build.wheel.packages]", pyproject) + self.assertIn('TileOps = "lib/TileOps"', pyproject) + self.assertTrue((project_root / "lib/TileOps/__init__.py").is_file()) if __name__ == "__main__": From 1f81e00d4e44b65e80db7ded293841dc7228a00f Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Mon, 10 Aug 2026 11:31:39 +0800 Subject: [PATCH 078/122] fix: address tfillpad review findings --- .github/workflows/ci.yml | 8 +++ docs/isa/tile-op/12-fill-and-padding-ops.md | 2 +- .../Transforms/InsertTemplateAttributes.cpp | 3 ++ lib/TileOps/a5/_fillpad.py | 26 ++++++++-- .../pto/textract_acc_to_vec_insert_sync.pto | 35 +++++++++++++ .../pto/tfillpad_non_normal_mat_invalid.pto | 8 +++ .../vpto/expand_tile_op_tilelang_tfillpad.pto | 12 +++-- test/samples/runop.sh | 7 ++- tools/ptobc/MAINTENANCE.md | 1 + tools/ptobc/generated/ptobc_opcodes_v0.h | 3 ++ .../testdata/tfillpad_legacy_v0_roundtrip.pto | 17 +++++++ tools/ptobc/tests/CMakeLists.txt | 9 ++++ .../ptobc/tests/fp_operand_forms_v0_encode.sh | 5 +- .../ptobc/tests/tfillpad_legacy_v0_decode.sh | 50 +++++++++++++++++++ .../tests/v0_fp_schema_compatibility_check.py | 3 ++ 15 files changed, 179 insertions(+), 10 deletions(-) create mode 100644 test/lit/pto/textract_acc_to_vec_insert_sync.pto create mode 100644 tools/ptobc/testdata/tfillpad_legacy_v0_roundtrip.pto create mode 100755 tools/ptobc/tests/tfillpad_legacy_v0_decode.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5cc540930a..102a7e65a8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -326,6 +326,14 @@ jobs: export PATH="${PTOAS_VENV}/bin:${LLVM_DIR}/bin:${PATH}" ctest --test-dir "${PTO_BUILD_DIR}" --output-on-failure -L PTODSL + - name: Run PTO-BC tests + shell: bash + env: + LLVM_DIR: ${{ env.LLVM_DIR }} + run: | + export PATH="${PTOAS_VENV}/bin:${LLVM_DIR}/bin:${PATH}" + ctest --test-dir "${PTO_BUILD_DIR}" --output-on-failure -R '^ptobc_' + - name: Run lit tests shell: bash env: diff --git a/docs/isa/tile-op/12-fill-and-padding-ops.md b/docs/isa/tile-op/12-fill-and-padding-ops.md index f38c6fe02b..4bbc77c604 100644 --- a/docs/isa/tile-op/12-fill-and-padding-ops.md +++ b/docs/isa/tile-op/12-fill-and-padding-ops.md @@ -41,7 +41,7 @@ pto.tfillpad ins(%src : !pto.tile_buf<...>) - The destination tile must carry a meaningful pad configuration. - In-place and expand lowering are VEC-only. Normal lowering also supports the homogeneous MAT overload. - Expand inference compares physical `shape`, not `valid_shape`. -- When physical shapes are equal, PTOAS compares exact starting addresses after PlanMemory. If equality cannot be proven, it conservatively chooses normal lowering. +- When physical shapes are equal, PTOAS compares exact starting addresses after PlanMemory. If equality cannot be proven, it conservatively chooses alias-safe normal lowering, which copies the complete valid region before writing padding. - MAT always uses Normal lowering, including when source and destination share the same starting address. **Example:** diff --git a/lib/PTO/Transforms/InsertTemplateAttributes.cpp b/lib/PTO/Transforms/InsertTemplateAttributes.cpp index 12016c8608..324d8c0a3f 100644 --- a/lib/PTO/Transforms/InsertTemplateAttributes.cpp +++ b/lib/PTO/Transforms/InsertTemplateAttributes.cpp @@ -473,6 +473,9 @@ static bool tryAppendPrecisionType( return true; } +// Candidate discovery runs before memory planning, so address-dependent +// context such as tfillpad's lowering_kind intentionally belongs only to the +// post-planning specialization key built by ExpandTileOp. static void appendOpContextAttrs( Operation *op, SmallVectorImpl> &attrs) { if (auto tcvt = dyn_cast(op)) { diff --git a/lib/TileOps/a5/_fillpad.py b/lib/TileOps/a5/_fillpad.py index d7db069fac..e8e7f1a7e7 100644 --- a/lib/TileOps/a5/_fillpad.py +++ b/lib/TileOps/a5/_fillpad.py @@ -211,12 +211,30 @@ def template(src: pto.Tile, dst: pto.Tile): if lowering_kind == "in_place": _fill_inplace(dst, src_valid_rows, src_valid_cols, dst_valid_rows, dst_valid_cols) return + if lowering_kind == "normal": + # Normal is the conservative fallback when memory planning cannot + # prove whether the two tile addresses alias. Preserve the entire + # valid source region before padding so this path is correct for + # both distinct and exactly aliased storage. + _copy_region(src, dst, src_valid_rows, 0, src_valid_cols) + _fill_inplace( + dst, + src_valid_rows, + src_valid_cols, + dst_valid_rows, + dst_valid_cols, + ) + return _copy_region(src, dst, src_valid_rows, 0, aligned_cols) - fill_row_stop = ( - dst_valid_rows if lowering_kind == "expand" else src_valid_rows - ) scalar_tail_start = _scalar_tail_start(dst, lanes) - _fill(dst, 0, fill_row_stop, aligned_cols, dst_valid_cols, scalar_tail_start=scalar_tail_start) + _fill( + dst, + 0, + dst_valid_rows, + aligned_cols, + dst_valid_cols, + scalar_tail_start=scalar_tail_start, + ) _copy_region(src, dst, src_valid_rows, aligned_cols, src_valid_cols) _fill(dst, src_valid_rows, dst_valid_rows, 0, dst_valid_cols, scalar_tail_start=scalar_tail_start) diff --git a/test/lit/pto/textract_acc_to_vec_insert_sync.pto b/test/lit/pto/textract_acc_to_vec_insert_sync.pto new file mode 100644 index 0000000000..52c5d403ad --- /dev/null +++ b/test/lit/pto/textract_acc_to_vec_insert_sync.pto @@ -0,0 +1,35 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --enable-insert-sync \ +// RUN: --emit-pto-ir %s -o - | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + func.func @textract_acc_to_vec_sync() { + %c0 = arith.constant 0 : index + %addr0 = arith.constant 0 : i64 + %addr1 = arith.constant 4096 : i64 + %addr2 = arith.constant 8192 : i64 + %src = pto.alloc_tile addr = %addr0 : !pto.tile_buf + %mid = pto.alloc_tile addr = %addr1 : !pto.tile_buf + %dst = pto.alloc_tile addr = %addr2 : !pto.tile_buf + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + outs(%mid : !pto.tile_buf) + {accToVecMode = #pto.acc_to_vec_mode} + pto.tabs ins(%mid : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK-LABEL: func.func @textract_acc_to_vec_sync +// CHECK: pto.textract +// CHECK: pto.set_flag[, , ] +// CHECK: pto.wait_flag[, , ] +// CHECK: pto.tabs diff --git a/test/lit/pto/tfillpad_non_normal_mat_invalid.pto b/test/lit/pto/tfillpad_non_normal_mat_invalid.pto index 9c42d41ec8..35abe95886 100644 --- a/test/lit/pto/tfillpad_non_normal_mat_invalid.pto +++ b/test/lit/pto/tfillpad_non_normal_mat_invalid.pto @@ -1,3 +1,11 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + // RUN: not ptoas --pto-arch=a3 %s -o /dev/null 2>&1 | FileCheck %s module { diff --git a/test/lit/vpto/expand_tile_op_tilelang_tfillpad.pto b/test/lit/vpto/expand_tile_op_tilelang_tfillpad.pto index 45fcf59a6f..2e0c89f927 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tfillpad.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tfillpad.pto @@ -19,9 +19,15 @@ // CHECK-NOT: pto.tfillpad ins // CHECK: pto.vecscope // CHECK: pto.castptr -// CHECK-DAG: pto.vdup -// CHECK-DAG: pto.vlds -// CHECK-DAG: pto.vsts +// Normal lowering is the conservative fallback for unprovable addresses. It +// must finish copying the valid source region before issuing any pad stores so +// the generated template remains correct when the addresses alias at runtime. +// CHECK-NOT: pto.vdup +// CHECK: pto.vlds +// CHECK: pto.vsts +// CHECK: pto.pxor +// CHECK: pto.vdup +// CHECK: pto.vsts // Note: vstus is not supported in TileLang DSL v1, so padding uses vsts instead module attributes {pto.kernel_kind = #pto.kernel_kind} { diff --git a/test/samples/runop.sh b/test/samples/runop.sh index d8258a79c6..db138d9b40 100755 --- a/test/samples/runop.sh +++ b/test/samples/runop.sh @@ -988,11 +988,16 @@ PY fi if [[ "$base" == "fillpad" ]]; then - if ! grep -Fq "TFILLPAD" "$cpp"; then + if ! grep -Fq "TFILLPAD(" "$cpp"; then echo -e "${A}(${base}.py)\tFAIL\tmissing compiler-inferred TFILLPAD lowering" overall=1 continue fi + if grep -Fq "TFILLPAD, + %dst: !pto.tile_buf) { + pto.tfillpad ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} diff --git a/tools/ptobc/tests/CMakeLists.txt b/tools/ptobc/tests/CMakeLists.txt index 2afc1f5d67..879b1390c7 100644 --- a/tools/ptobc/tests/CMakeLists.txt +++ b/tools/ptobc/tests/CMakeLists.txt @@ -42,6 +42,15 @@ add_test(NAME ptobc_v0_fp_schema_compatibility_check ${CMAKE_SOURCE_DIR}/tools/ptobc/generated/ptobc_opcodes_v0.h ) +add_test(NAME ptobc_tfillpad_legacy_v0_decode + COMMAND ${CMAKE_COMMAND} -E env + PTOBC_BIN=$ + PTOAS_BIN=${CMAKE_BINARY_DIR}/tools/ptoas/ptoas + PYTHON_EXECUTABLE=${Python3_EXECUTABLE} + TESTDATA_DIR=${PTObc_TESTDATA_DIR} + ${CMAKE_CURRENT_LIST_DIR}/tfillpad_legacy_v0_decode.sh +) + add_test(NAME ptobc_trowexpandsub_v0_encode COMMAND ${CMAKE_COMMAND} -E env PTOBC_BIN=$ diff --git a/tools/ptobc/tests/fp_operand_forms_v0_encode.sh b/tools/ptobc/tests/fp_operand_forms_v0_encode.sh index c29d0c98c7..6b2e228f36 100755 --- a/tools/ptobc/tests/fp_operand_forms_v0_encode.sh +++ b/tools/ptobc/tests/fp_operand_forms_v0_encode.sh @@ -46,7 +46,10 @@ PY grep -F "pto.textract ins(" "${ROUNDTRIP}" >/dev/null grep -F "pto.tinsert ins(" "${ROUNDTRIP}" >/dev/null grep -F "pto.tmov ins(" "${ROUNDTRIP}" >/dev/null -[[ $(grep -Fc " fp " "${ROUNDTRIP}") -eq 3 ]] +[[ $(grep -Fc " fp " "${ROUNDTRIP}") -eq 2 ]] +grep -F "pto.tmov ins(" "${ROUNDTRIP}" \ + | grep -F ", %" \ + | grep -F "!pto.tile_buf/dev/null # Parsing and verification prove that decoder-reconstructed segment metadata is # valid, including records encoded with the legacy FP wire operand ordering. diff --git a/tools/ptobc/tests/tfillpad_legacy_v0_decode.sh b/tools/ptobc/tests/tfillpad_legacy_v0_decode.sh new file mode 100755 index 0000000000..b86b44d1e2 --- /dev/null +++ b/tools/ptobc/tests/tfillpad_legacy_v0_decode.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +set -euo pipefail + +PTOBC_BIN=${PTOBC_BIN:-} +PTOAS_BIN=${PTOAS_BIN:-} +TESTDATA_DIR=${TESTDATA_DIR:-} +PYTHON_EXECUTABLE=${PYTHON_EXECUTABLE:-} +if [[ -z "${PTOBC_BIN}" || -z "${PTOAS_BIN}" || -z "${PYTHON_EXECUTABLE}" || -z "${TESTDATA_DIR}" ]]; then + echo "error: PTOBC_BIN, PTOAS_BIN, PYTHON_EXECUTABLE, and TESTDATA_DIR must be set" >&2 + exit 2 +fi + +IN="${TESTDATA_DIR}/tfillpad_legacy_v0_roundtrip.pto" +OUT_DIR=${OUT_DIR:-"${PWD}/ptobc_tfillpad_legacy_out"} +mkdir -p "${OUT_DIR}" + +CANONICAL="${OUT_DIR}/tfillpad_canonical.ptobc" +EXPAND="${OUT_DIR}/tfillpad_expand_legacy.ptobc" +INPLACE="${OUT_DIR}/tfillpad_inplace_legacy.ptobc" + +"${PTOBC_BIN}" encode "${IN}" -o "${CANONICAL}" + +"${PYTHON_EXECUTABLE}" - <<'PY' "${CANONICAL}" "${EXPAND}" "${INPLACE}" +from pathlib import Path +import sys + +canonical = Path(sys.argv[1]).read_bytes() +wire_opcode = b"\x23\x10" +if canonical.count(wire_opcode) != 1: + raise SystemExit("expected exactly one canonical tfillpad opcode") + +Path(sys.argv[2]).write_bytes(canonical.replace(wire_opcode, b"\x24\x10", 1)) +Path(sys.argv[3]).write_bytes(canonical.replace(wire_opcode, b"\x25\x10", 1)) +PY + +for kind in expand inplace; do + bc="${OUT_DIR}/tfillpad_${kind}_legacy.ptobc" + roundtrip="${OUT_DIR}/tfillpad_${kind}_legacy.roundtrip.pto" + "${PTOBC_BIN}" decode "${bc}" -o "${roundtrip}" + grep -F "pto.tfillpad ins(" "${roundtrip}" >/dev/null + "${PTOAS_BIN}" --pto-arch=a5 --emit-pto-ir "${roundtrip}" -o /dev/null +done diff --git a/tools/ptobc/tests/v0_fp_schema_compatibility_check.py b/tools/ptobc/tests/v0_fp_schema_compatibility_check.py index e8c5b8d99a..f6dea29196 100755 --- a/tools/ptobc/tests/v0_fp_schema_compatibility_check.py +++ b/tools/ptobc/tests/v0_fp_schema_compatibility_check.py @@ -15,6 +15,9 @@ EXPECTED = { 0x1021: ("pto.textract", 0x00, 4), 0x1022: ("pto.textract", 0x00, 5), + 0x1023: ("pto.tfillpad", 0x00, 2), + 0x1024: ("pto.tfillpad", 0x00, 2), + 0x1025: ("pto.tfillpad", 0x00, 2), 0x102D: ("pto.tinsert", 0x00, 4), 0x102E: ("pto.tinsert", 0x00, 5), 0x1038: ("pto.tmov", 0x00, 2), From 1de45a7e341bf2fd4985d69b51d7511bb2edb5e4 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Fri, 7 Aug 2026 17:11:24 +0800 Subject: [PATCH 079/122] test(vpto): cover vscatter CSE memory effects --- .../vscatter_cse_memory_effects_vpto_llvm.pto | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 test/lit/vpto/vscatter_cse_memory_effects_vpto_llvm.pto diff --git a/test/lit/vpto/vscatter_cse_memory_effects_vpto_llvm.pto b/test/lit/vpto/vscatter_cse_memory_effects_vpto_llvm.pto new file mode 100644 index 0000000000..bcb01274ae --- /dev/null +++ b/test/lit/vpto/vscatter_cse_memory_effects_vpto_llvm.pto @@ -0,0 +1,45 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software; you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// Guards the vscatter memory effects required to prevent Bisheng EarlyCSE from +// reusing a vlds result across an in-place scatter. +// RUN: ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-backend=vpto --emit-vpto-llvm-ir %s -o - 2>&1 | FileCheck %s +// RUN: ptoas --cann-output-version=9.0.0-beta.1 --pto-arch=a5 --pto-backend=vpto --emit-vpto-llvm-ir %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vscatter_cse_memory_effects(%data: !pto.ptr, + %out: !pto.ptr, + %value: !pto.vreg<64xf32>, + %offsets: !pto.vreg<64xi32>) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c64 = arith.constant 64 : index + pto.vecscope { + %mask = pto.pset_b32 "PAT_ALL" : !pto.mask + %before = pto.vlds %data[%c0] : !pto.ptr -> !pto.vreg<64xf32> + pto.vscatter %value, %data, %offsets, %mask : !pto.vreg<64xf32>, !pto.ptr, !pto.vreg<64xi32>, !pto.mask + pto.mem_bar "VST_VLD" + %after = pto.vlds %data[%c0] : !pto.ptr -> !pto.vreg<64xf32> + pto.vsts %before, %out[%c0], %mask : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.vsts %after, %out[%c64], %mask : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + } + return + } +} + +// CHECK: declare <64 x float> @llvm.hivm.vldsx1.v64f32 +// CHECK: declare void @llvm.hivm.vscatter.v64f32.v300({{.*}}) #[[VSCATTER_ATTR:[0-9]+]] +// CHECK: declare void @llvm.hivm.mem.bar.vst.vld() +// CHECK: declare void @llvm.hivm.vstsx1.v64f32 +// CHECK-LABEL: define void @vscatter_cse_memory_effects_mix_aiv +// CHECK: call <64 x float> @llvm.hivm.vldsx1.v64f32 +// CHECK: call void @llvm.hivm.vscatter.v64f32.v300 +// CHECK-NEXT: call void @llvm.hivm.mem.bar.vst.vld() +// CHECK-NEXT: {{.*}}call <64 x float> @llvm.hivm.vldsx1.v64f32 +// CHECK: call void @llvm.hivm.vstsx1.v64f32 +// CHECK: call void @llvm.hivm.vstsx1.v64f32 +// CHECK: attributes #[[VSCATTER_ATTR]] = { nounwind writeonly argmemonly From 5a06bd4a8489076d2d11905178c1535fdb75c625 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Fri, 7 Aug 2026 18:16:44 +0800 Subject: [PATCH 080/122] fix(vpto): attach vscatter memory effects at declaration --- lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp | 40 +++++++++++-------- lib/PTO/Transforms/VPTOLLVMEmitter.cpp | 40 +++++++++++-------- .../Transforms/VPTOLLVMEmitterDispatcher.cpp | 13 +++++- 3 files changed, 57 insertions(+), 36 deletions(-) diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index 2c60df94b9..0013383435 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -291,7 +291,7 @@ getVPTOStructFieldAddress(ConversionPatternRewriter &rewriter, Location loc, struct PlannedDecl { std::string name; FunctionType type; - bool writeOnlyDestination = false; + bool needsVscatterMemoryEffectsWorkaround = false; }; struct LoweringState { @@ -4442,17 +4442,33 @@ materializeDecls(ModuleOp module, ArrayRef plannedDecls, OpBuilder builder(module.getBodyRegion()); builder.setInsertionPointToStart(&module.getBodyRegion().front()); for (const PlannedDecl &decl : plannedDecls) { - if (func::FuncOp existing = module.lookupSymbol(decl.name)) { - if (existing.getFunctionType() != decl.type) { + func::FuncOp func = module.lookupSymbol(decl.name); + if (func) { + if (func.getFunctionType() != decl.type) { diagOS << "VPTO LLVM emission failed: conflicting declaration for " << decl.name << "\n"; return failure(); } - continue; + } else { + func = builder.create(module.getLoc(), decl.name, decl.type); + func.setPrivate(); } - auto func = - builder.create(module.getLoc(), decl.name, decl.type); - func.setPrivate(); + + if (!decl.needsVscatterMemoryEffectsWorkaround) + continue; + + // Work around a bug in older Bisheng releases: vscatter was not modeled + // as writing through its destination pointer, so EarlyCSE could eliminate + // a load after vscatter as redundant. Carry the intrinsic's real memory + // effects on its declaration through func-to-llvm. The dispatcher later + // rewrites LLVM 21's memory(...) spelling for Bisheng's LLVM 15 parser. + func->setAttr( + "memory_effects", + LLVM::MemoryEffectsAttr::get( + module.getContext(), + {LLVM::ModRefInfo::NoModRef, LLVM::ModRefInfo::Mod, + LLVM::ModRefInfo::NoModRef})); + func->setAttr("no_unwind", builder.getUnitAttr()); } return success(); } @@ -11551,16 +11567,6 @@ emitDeviceLLVMModule(ModuleOp deviceModule, StringRef kernelKind, } applyArtifactVisibilityLinkage(deviceModule, *llvmModule); - for (llvm::Function &func : *llvmModule) { - if (!func.getName().starts_with("llvm.hivm.vscatter.")) - continue; - // Bisheng LLVM 15 verifies these intrinsic memory effects. Record them - // through the LLVM 21 API here; the dispatcher rewrites the new textual - // memory(...) spelling before handing the IR to Bisheng. - func.setOnlyAccessesArgMemory(); - func.addFnAttr(llvm::Attribute::NoUnwind); - func.addFnAttr(llvm::Attribute::WriteOnly); - } applySimtEntryCallingConvention(*llvmModule, simtEntryNames); if (failed(attachAIVectorScopeMetadata(*llvmModule, diagOS))) return failure(); diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 81a3ac9bd7..7af367d836 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -293,7 +293,7 @@ getVPTOStructFieldAddress(ConversionPatternRewriter &rewriter, Location loc, struct PlannedDecl { std::string name; FunctionType type; - bool writeOnlyDestination = false; + bool needsVscatterMemoryEffectsWorkaround = false; }; struct LoweringState { @@ -4487,17 +4487,33 @@ materializeDecls(ModuleOp module, ArrayRef plannedDecls, OpBuilder builder(module.getBodyRegion()); builder.setInsertionPointToStart(&module.getBodyRegion().front()); for (const PlannedDecl &decl : plannedDecls) { - if (func::FuncOp existing = module.lookupSymbol(decl.name)) { - if (existing.getFunctionType() != decl.type) { + func::FuncOp func = module.lookupSymbol(decl.name); + if (func) { + if (func.getFunctionType() != decl.type) { diagOS << "VPTO LLVM emission failed: conflicting declaration for " << decl.name << "\n"; return failure(); } - continue; + } else { + func = builder.create(module.getLoc(), decl.name, decl.type); + func.setPrivate(); } - auto func = - builder.create(module.getLoc(), decl.name, decl.type); - func.setPrivate(); + + if (!decl.needsVscatterMemoryEffectsWorkaround) + continue; + + // Work around a bug in older Bisheng releases: vscatter was not modeled + // as writing through its destination pointer, so EarlyCSE could eliminate + // a load after vscatter as redundant. Carry the intrinsic's real memory + // effects on its declaration through func-to-llvm. The dispatcher later + // rewrites LLVM 21's memory(...) spelling for Bisheng's LLVM 15 parser. + func->setAttr( + "memory_effects", + LLVM::MemoryEffectsAttr::get( + module.getContext(), + {LLVM::ModRefInfo::NoModRef, LLVM::ModRefInfo::Mod, + LLVM::ModRefInfo::NoModRef})); + func->setAttr("no_unwind", builder.getUnitAttr()); } return success(); } @@ -12249,16 +12265,6 @@ emitDeviceLLVMModule(ModuleOp deviceModule, StringRef kernelKind, } applyArtifactVisibilityLinkage(deviceModule, *llvmModule); - for (llvm::Function &func : *llvmModule) { - if (!func.getName().starts_with("llvm.hivm.vscatter.")) - continue; - // Bisheng LLVM 15 verifies these intrinsic memory effects. Record them - // through the LLVM 21 API here; the dispatcher rewrites the new textual - // memory(...) spelling before handing the IR to Bisheng. - func.setOnlyAccessesArgMemory(); - func.addFnAttr(llvm::Attribute::NoUnwind); - func.addFnAttr(llvm::Attribute::WriteOnly); - } applySimtEntryCallingConvention(*llvmModule, simtEntryNames); if (failed(attachAIVectorScopeMetadata(*llvmModule, diagOS))) return failure(); diff --git a/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp b/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp index 70c3b852eb..adee1a1ed0 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp @@ -85,9 +85,18 @@ LogicalResult lowerVPTOModuleToLLVMIRText( } os.flush(); // LLVM 21 prints arg-memory effects with memory(...), while the Bisheng - // LLVM 15 parser accepts only the equivalent legacy argmemonly spelling. - constexpr StringLiteral modernArgMemOnly = "memory(argmem: readwrite)"; + // LLVM 15 parser accepts only the equivalent legacy spellings. + constexpr StringLiteral modernWriteOnlyArgMem = "memory(argmem: write)"; size_t offset = 0; + while ((offset = output.find(modernWriteOnlyArgMem.str(), offset)) != + std::string::npos) { + output.replace(offset, modernWriteOnlyArgMem.size(), + "writeonly argmemonly"); + offset += StringRef("writeonly argmemonly").size(); + } + + constexpr StringLiteral modernArgMemOnly = "memory(argmem: readwrite)"; + offset = 0; while ((offset = output.find(modernArgMemOnly.str(), offset)) != std::string::npos) { output.replace(offset, modernArgMemOnly.size(), "argmemonly"); From 87106949a28ea7aa8278ebc960b5f09ca1709260 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Mon, 10 Aug 2026 10:55:45 +0800 Subject: [PATCH 081/122] fix(vpto): keep vscatter Bisheng workaround localized --- lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp | 42 ++++++++----------- lib/PTO/Transforms/VPTOLLVMEmitter.cpp | 42 ++++++++----------- .../Transforms/VPTOLLVMEmitterDispatcher.cpp | 13 +----- 3 files changed, 36 insertions(+), 61 deletions(-) diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index 0013383435..767b409633 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -291,7 +291,6 @@ getVPTOStructFieldAddress(ConversionPatternRewriter &rewriter, Location loc, struct PlannedDecl { std::string name; FunctionType type; - bool needsVscatterMemoryEffectsWorkaround = false; }; struct LoweringState { @@ -4442,33 +4441,17 @@ materializeDecls(ModuleOp module, ArrayRef plannedDecls, OpBuilder builder(module.getBodyRegion()); builder.setInsertionPointToStart(&module.getBodyRegion().front()); for (const PlannedDecl &decl : plannedDecls) { - func::FuncOp func = module.lookupSymbol(decl.name); - if (func) { - if (func.getFunctionType() != decl.type) { + if (func::FuncOp existing = module.lookupSymbol(decl.name)) { + if (existing.getFunctionType() != decl.type) { diagOS << "VPTO LLVM emission failed: conflicting declaration for " << decl.name << "\n"; return failure(); } - } else { - func = builder.create(module.getLoc(), decl.name, decl.type); - func.setPrivate(); - } - - if (!decl.needsVscatterMemoryEffectsWorkaround) continue; - - // Work around a bug in older Bisheng releases: vscatter was not modeled - // as writing through its destination pointer, so EarlyCSE could eliminate - // a load after vscatter as redundant. Carry the intrinsic's real memory - // effects on its declaration through func-to-llvm. The dispatcher later - // rewrites LLVM 21's memory(...) spelling for Bisheng's LLVM 15 parser. - func->setAttr( - "memory_effects", - LLVM::MemoryEffectsAttr::get( - module.getContext(), - {LLVM::ModRefInfo::NoModRef, LLVM::ModRefInfo::Mod, - LLVM::ModRefInfo::NoModRef})); - func->setAttr("no_unwind", builder.getUnitAttr()); + } + auto func = + builder.create(module.getLoc(), decl.name, decl.type); + func.setPrivate(); } return success(); } @@ -7857,8 +7840,7 @@ class LowerVscatterOpPattern final op.getLoc(), *calleeName, TypeRange{}, ValueRange{adaptor.getValue(), adaptor.getDestination(), adaptor.getOffsets(), adaptor.getMask()}); - state.plannedDecls.push_back( - PlannedDecl{calleeName->str(), funcType, true}); + state.plannedDecls.push_back(PlannedDecl{calleeName->str(), funcType}); rewriter.eraseOp(op); return success(); } @@ -11567,6 +11549,16 @@ emitDeviceLLVMModule(ModuleOp deviceModule, StringRef kernelKind, } applyArtifactVisibilityLinkage(deviceModule, *llvmModule); + for (llvm::Function &func : *llvmModule) { + if (!func.getName().starts_with("llvm.hivm.vscatter.")) + continue; + // Work around a bug in older Bisheng releases: vscatter was not modeled + // as writing through its destination pointer, so EarlyCSE could eliminate + // a load after vscatter as redundant. + func.setOnlyAccessesArgMemory(); + func.addFnAttr(llvm::Attribute::NoUnwind); + func.addFnAttr(llvm::Attribute::WriteOnly); + } applySimtEntryCallingConvention(*llvmModule, simtEntryNames); if (failed(attachAIVectorScopeMetadata(*llvmModule, diagOS))) return failure(); diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 7af367d836..e98e091dee 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -293,7 +293,6 @@ getVPTOStructFieldAddress(ConversionPatternRewriter &rewriter, Location loc, struct PlannedDecl { std::string name; FunctionType type; - bool needsVscatterMemoryEffectsWorkaround = false; }; struct LoweringState { @@ -4487,33 +4486,17 @@ materializeDecls(ModuleOp module, ArrayRef plannedDecls, OpBuilder builder(module.getBodyRegion()); builder.setInsertionPointToStart(&module.getBodyRegion().front()); for (const PlannedDecl &decl : plannedDecls) { - func::FuncOp func = module.lookupSymbol(decl.name); - if (func) { - if (func.getFunctionType() != decl.type) { + if (func::FuncOp existing = module.lookupSymbol(decl.name)) { + if (existing.getFunctionType() != decl.type) { diagOS << "VPTO LLVM emission failed: conflicting declaration for " << decl.name << "\n"; return failure(); } - } else { - func = builder.create(module.getLoc(), decl.name, decl.type); - func.setPrivate(); - } - - if (!decl.needsVscatterMemoryEffectsWorkaround) continue; - - // Work around a bug in older Bisheng releases: vscatter was not modeled - // as writing through its destination pointer, so EarlyCSE could eliminate - // a load after vscatter as redundant. Carry the intrinsic's real memory - // effects on its declaration through func-to-llvm. The dispatcher later - // rewrites LLVM 21's memory(...) spelling for Bisheng's LLVM 15 parser. - func->setAttr( - "memory_effects", - LLVM::MemoryEffectsAttr::get( - module.getContext(), - {LLVM::ModRefInfo::NoModRef, LLVM::ModRefInfo::Mod, - LLVM::ModRefInfo::NoModRef})); - func->setAttr("no_unwind", builder.getUnitAttr()); + } + auto func = + builder.create(module.getLoc(), decl.name, decl.type); + func.setPrivate(); } return success(); } @@ -8456,8 +8439,7 @@ class LowerVscatterOpPattern final op.getLoc(), *calleeName, TypeRange{}, ValueRange{adaptor.getValue(), adaptor.getDestination(), adaptor.getOffsets(), adaptor.getMask()}); - state.plannedDecls.push_back( - PlannedDecl{calleeName->str(), funcType, true}); + state.plannedDecls.push_back(PlannedDecl{calleeName->str(), funcType}); rewriter.eraseOp(op); return success(); } @@ -12265,6 +12247,16 @@ emitDeviceLLVMModule(ModuleOp deviceModule, StringRef kernelKind, } applyArtifactVisibilityLinkage(deviceModule, *llvmModule); + for (llvm::Function &func : *llvmModule) { + if (!func.getName().starts_with("llvm.hivm.vscatter.")) + continue; + // Work around a bug in older Bisheng releases: vscatter was not modeled + // as writing through its destination pointer, so EarlyCSE could eliminate + // a load after vscatter as redundant. + func.setOnlyAccessesArgMemory(); + func.addFnAttr(llvm::Attribute::NoUnwind); + func.addFnAttr(llvm::Attribute::WriteOnly); + } applySimtEntryCallingConvention(*llvmModule, simtEntryNames); if (failed(attachAIVectorScopeMetadata(*llvmModule, diagOS))) return failure(); diff --git a/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp b/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp index adee1a1ed0..70c3b852eb 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp @@ -85,18 +85,9 @@ LogicalResult lowerVPTOModuleToLLVMIRText( } os.flush(); // LLVM 21 prints arg-memory effects with memory(...), while the Bisheng - // LLVM 15 parser accepts only the equivalent legacy spellings. - constexpr StringLiteral modernWriteOnlyArgMem = "memory(argmem: write)"; - size_t offset = 0; - while ((offset = output.find(modernWriteOnlyArgMem.str(), offset)) != - std::string::npos) { - output.replace(offset, modernWriteOnlyArgMem.size(), - "writeonly argmemonly"); - offset += StringRef("writeonly argmemonly").size(); - } - + // LLVM 15 parser accepts only the equivalent legacy argmemonly spelling. constexpr StringLiteral modernArgMemOnly = "memory(argmem: readwrite)"; - offset = 0; + size_t offset = 0; while ((offset = output.find(modernArgMemOnly.str(), offset)) != std::string::npos) { output.replace(offset, modernArgMemOnly.size(), "argmemonly"); From 544da0adcfe6b7b659baa0fe416b3ede33790483 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Mon, 10 Aug 2026 15:56:47 +0800 Subject: [PATCH 082/122] fix: preserve extended FP v0 semantics --- .github/workflows/ci.yml | 42 +++++++++ tools/ptobc/MAINTENANCE.md | 6 ++ tools/ptobc/src/mlir_encode.cpp | 48 ++++++++-- .../tmov_fp_extended_v0_roundtrip.pto | 21 +++++ .../testdata/tstore_fp_extended_v0_reject.pto | 25 +++++ tools/ptobc/tests/CMakeLists.txt | 9 ++ .../tests/fp_extended_v0_compatibility.sh | 94 +++++++++++++++++++ tools/ptobc/tests/stage9_e2e.sh | 1 + 8 files changed, 238 insertions(+), 8 deletions(-) create mode 100644 tools/ptobc/testdata/tmov_fp_extended_v0_roundtrip.pto create mode 100644 tools/ptobc/testdata/tstore_fp_extended_v0_reject.pto create mode 100755 tools/ptobc/tests/fp_extended_v0_compatibility.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 102a7e65a8..377850945e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -334,6 +334,48 @@ jobs: export PATH="${PTOAS_VENV}/bin:${LLVM_DIR}/bin:${PATH}" ctest --test-dir "${PTO_BUILD_DIR}" --output-on-failure -R '^ptobc_' + - name: Verify PTO-BC v0 against the pre-unification reader + shell: bash + env: + PTOBC_V0_LEGACY_COMMIT: 9c49c3697de35d4b36e2abc5a00da0b264ae1bb6 + run: | + set -euo pipefail + legacy_root="${RUNNER_TEMP:-${GITHUB_WORKSPACE}/.tmp}/ptoas-v0-legacy" + legacy_src="${legacy_root}/src" + legacy_build="${legacy_root}/build" + rm -rf "${legacy_root}" + git fetch --no-tags --depth=1 \ + https://github.com/hw-native-sys/PTOAS.git \ + "${PTOBC_V0_LEGACY_COMMIT}" + git worktree add --detach "${legacy_src}" FETCH_HEAD + + cmake -C "${legacy_src}/cmake/LinuxHardeningCache.cmake" \ + -G Ninja \ + -S "${legacy_src}" \ + -B "${legacy_build}" \ + -DLLVM_DIR="${LLVM_DIR}/lib/cmake/llvm" \ + -DMLIR_DIR="${LLVM_DIR}/lib/cmake/mlir" \ + -DPython3_EXECUTABLE="${PTOAS_VENV}/bin/python" \ + -DPython_EXECUTABLE="${PTOAS_VENV}/bin/python" \ + -DPython3_FIND_STRATEGY=LOCATION \ + -Dpybind11_DIR="$("${PTOAS_VENV}/bin/python" -m pybind11 --cmakedir)" \ + -Dnanobind_DIR="$("${PTOAS_VENV}/bin/python" -m nanobind --cmake_dir)" \ + -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ + -DPTO_ENABLE_PYTHON_BINDING=ON \ + -DBUILD_TESTING=OFF \ + -DCMAKE_C_COMPILER="${PTOAS_CMAKE_C_COMPILER}" \ + -DCMAKE_CXX_COMPILER="${PTOAS_CMAKE_CXX_COMPILER}" + cmake --build "${legacy_build}" --target PTOASPythonPackage ptobc + + PTOBC_BIN="${PTO_BUILD_DIR}/tools/ptobc/ptobc" \ + PTOAS_BIN="${PTO_BUILD_DIR}/tools/ptoas/ptoas" \ + PYTHON_EXECUTABLE="${PTOAS_VENV}/bin/python" \ + TESTDATA_DIR="${GITHUB_WORKSPACE}/tools/ptobc/testdata" \ + LEGACY_PTOBC_BIN="${legacy_build}/tools/ptobc/ptobc" \ + LEGACY_PTOAS_BIN="${legacy_build}/tools/ptoas/ptoas" \ + OUT_DIR="${legacy_root}/cross-version-output" \ + bash tools/ptobc/tests/fp_extended_v0_compatibility.sh + - name: Run lit tests shell: bash env: diff --git a/tools/ptobc/MAINTENANCE.md b/tools/ptobc/MAINTENANCE.md index 1dec1b5fd7..32056e719c 100644 --- a/tools/ptobc/MAINTENANCE.md +++ b/tools/ptobc/MAINTENANCE.md @@ -30,6 +30,12 @@ Run (or rely on CI): - `ctest -R ptobc_opcode_coverage_check` - `ctest -R ptobc_v0_fp_schema_compatibility_check` - `ctest -R ptobc_tfillpad_legacy_v0_decode` +- `ctest -R ptobc_fp_extended_v0_compatibility` + +CI additionally builds the last pre-unification v0 reader at commit +`9c49c3697de35d4b36e2abc5a00da0b264ae1bb6` and runs the extended FP test as +new writer -> legacy reader -> legacy lowering. Update that pin only when the +minimum supported v0 reader changes deliberately. ## Notes - `ptobc_opcode_coverage_check` is a heuristic based on `mnemonic = "..."` occurrences. diff --git a/tools/ptobc/src/mlir_encode.cpp b/tools/ptobc/src/mlir_encode.cpp index 232a6cf23f..d96dc129af 100644 --- a/tools/ptobc/src/mlir_encode.cpp +++ b/tools/ptobc/src/mlir_encode.cpp @@ -74,6 +74,16 @@ using FunctionVector = llvm::SmallVector(textract.getPreQuantScalar()); if (auto tinsert = llvm::dyn_cast(&op)) return static_cast(tinsert.getPreQuantScalar()); - if (auto tmov = llvm::dyn_cast(&op)) - return static_cast(tmov.getPreQuantScalar()); + if (auto tmov = llvm::dyn_cast(&op)) { + if (tmov.getPreQuantScalar()) + return true; + // The removed pto.tmov.fp op carried only src/fp/dst. Its legacy opcode + // would silently discard mode/relu semantics when read by an older PTOAS. + // The pre-unification pto.tmov op already understood fp plus these attrs, + // so use the generic v0 record for the extended unified form. + return tmov.getFp() && !canUseLegacyTMovFpWireOpcode(tmov); + } if (llvm::isa( &op)) return true; @@ -109,8 +126,7 @@ static bool shouldEncodeViaGenericV0CompatibilityShim(mlir::Operation &op) { return false; } -static std::optional -getLegacyFpWireOpcode(mlir::Operation &op) { +static std::optional getLegacyFpWireOpcode(mlir::Operation &op) { if (auto textract = llvm::dyn_cast(&op)) return textract.getFp() ? std::optional(kTExtractFpWireOpcode) : std::nullopt; @@ -118,14 +134,27 @@ getLegacyFpWireOpcode(mlir::Operation &op) { return tinsert.getFp() ? std::optional(kTInsertFpWireOpcode) : std::nullopt; if (auto tmov = llvm::dyn_cast(&op)) - return tmov.getFp() ? std::optional(kTMovFpWireOpcode) - : std::nullopt; + return tmov.getFp() && canUseLegacyTMovFpWireOpcode(tmov) + ? std::optional(kTMovFpWireOpcode) + : std::nullopt; if (auto tstore = llvm::dyn_cast(&op)) - return tstore.getFp() ? std::optional(kTStoreFpWireOpcode) - : std::nullopt; + return tstore.getFp() && canUseLegacyTStoreFpWireOpcode(tstore) + ? std::optional(kTStoreFpWireOpcode) + : std::nullopt; return std::nullopt; } +static std::optional +getUnsupportedV0EncodingReason(mlir::Operation &op) { + auto tstore = llvm::dyn_cast(&op); + if (!tstore || !tstore.getFp() || canUseLegacyTStoreFpWireOpcode(tstore)) + return std::nullopt; + + return "pto.tstore fp with non-default atomicType or reluPreMode cannot " + "be represented safely in PTO-BC v0; legacy opcode 0x1066 would " + "silently drop those semantics"; +} + static bool omitsDerivedOperandSegmentsInV0(uint16_t opcode) { switch (opcode) { case kTExtractOpcode: @@ -756,6 +785,9 @@ void Encoder::encodeOp(mlir::Operation& op, Buffer& out) { } auto fullName = op.getName().getStringRef(); + if (auto reason = getUnsupportedV0EncodingReason(op)) + throw std::runtime_error(*reason); + if (auto tscatter = llvm::dyn_cast(&op)) { uint16_t opcode = tscatter.getMaskPatternAttr() ? ptobc::v0::kTscatterMaskOpcode diff --git a/tools/ptobc/testdata/tmov_fp_extended_v0_roundtrip.pto b/tools/ptobc/testdata/tmov_fp_extended_v0_roundtrip.pto new file mode 100644 index 0000000000..6c924891a4 --- /dev/null +++ b/tools/ptobc/testdata/tmov_fp_extended_v0_roundtrip.pto @@ -0,0 +1,21 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +module attributes {"pto.device-spec" = "Ascend950"} { + func.func @tmov_fp_extended_v0() { + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%src : !pto.tile_buf, + %fp : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {accToVecMode = #pto.acc_to_vec_mode, + reluPreMode = #pto} + return + } +} diff --git a/tools/ptobc/testdata/tstore_fp_extended_v0_reject.pto b/tools/ptobc/testdata/tstore_fp_extended_v0_reject.pto new file mode 100644 index 0000000000..7127eb97a5 --- /dev/null +++ b/tools/ptobc/testdata/tstore_fp_extended_v0_reject.pto @@ -0,0 +1,25 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +module attributes {"pto.device-spec" = "Ascend950"} { + func.func @tstore_fp_extended_v0_reject(%dst: !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 32 : index + %dst_tv = pto.make_tensor_view %dst, shape = [%c32, %c32], strides = [%c32, %c1] : !pto.tensor_view<32x32xi8> + %dst_part = pto.partition_view %dst_tv, offsets = [%c0, %c0], sizes = [%c32, %c32] : !pto.tensor_view<32x32xi8> -> !pto.partition_tensor_view<32x32xi8> + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + pto.tstore ins(%src : !pto.tile_buf + fp %fp : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<32x32xi8>) + {atomicType = #pto, + reluPreMode = #pto} + return + } +} diff --git a/tools/ptobc/tests/CMakeLists.txt b/tools/ptobc/tests/CMakeLists.txt index 879b1390c7..ad60459fe5 100644 --- a/tools/ptobc/tests/CMakeLists.txt +++ b/tools/ptobc/tests/CMakeLists.txt @@ -139,6 +139,15 @@ add_test(NAME ptobc_fp_operand_forms_v0_encode ${CMAKE_CURRENT_LIST_DIR}/fp_operand_forms_v0_encode.sh ) +add_test(NAME ptobc_fp_extended_v0_compatibility + COMMAND ${CMAKE_COMMAND} -E env + PTOBC_BIN=$ + PTOAS_BIN=${CMAKE_BINARY_DIR}/tools/ptoas/ptoas + PYTHON_EXECUTABLE=${Python3_EXECUTABLE} + TESTDATA_DIR=${PTObc_TESTDATA_DIR} + ${CMAKE_CURRENT_LIST_DIR}/fp_extended_v0_compatibility.sh +) + add_test(NAME ptobc_comm_p2p_dynamic_v0_encode COMMAND ${CMAKE_COMMAND} -E env PTOBC_BIN=$ diff --git a/tools/ptobc/tests/fp_extended_v0_compatibility.sh b/tools/ptobc/tests/fp_extended_v0_compatibility.sh new file mode 100755 index 0000000000..108a719d83 --- /dev/null +++ b/tools/ptobc/tests/fp_extended_v0_compatibility.sh @@ -0,0 +1,94 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +set -euo pipefail + +PTOBC_BIN=${PTOBC_BIN:-} +PTOAS_BIN=${PTOAS_BIN:-} +PYTHON_EXECUTABLE=${PYTHON_EXECUTABLE:-} +TESTDATA_DIR=${TESTDATA_DIR:-} +LEGACY_PTOBC_BIN=${LEGACY_PTOBC_BIN:-} +LEGACY_PTOAS_BIN=${LEGACY_PTOAS_BIN:-} +if [[ -z "${PTOBC_BIN}" || -z "${PTOAS_BIN}" || -z "${PYTHON_EXECUTABLE}" || -z "${TESTDATA_DIR}" ]]; then + echo "error: PTOBC_BIN, PTOAS_BIN, PYTHON_EXECUTABLE, and TESTDATA_DIR must be set" >&2 + exit 2 +fi +if [[ -n "${LEGACY_PTOBC_BIN}" && -z "${LEGACY_PTOAS_BIN}" ]] || + [[ -z "${LEGACY_PTOBC_BIN}" && -n "${LEGACY_PTOAS_BIN}" ]]; then + echo "error: LEGACY_PTOBC_BIN and LEGACY_PTOAS_BIN must be set together" >&2 + exit 2 +fi + +OUT_DIR=${OUT_DIR:-"${PWD}/ptobc_fp_extended_v0_out"} +mkdir -p "${OUT_DIR}" + +TMOV_IN="${TESTDATA_DIR}/tmov_fp_extended_v0_roundtrip.pto" +TMOV_BC="${OUT_DIR}/tmov_fp_extended_v0_roundtrip.ptobc" +TMOV_CURRENT_IR="${OUT_DIR}/tmov_fp_extended_v0.current.pto" +TMOV_CURRENT_CPP="${OUT_DIR}/tmov_fp_extended_v0.current.cpp" + +"${PTOBC_BIN}" encode "${TMOV_IN}" -o "${TMOV_BC}" +"${PYTHON_EXECUTABLE}" - <<'PY' "${TMOV_BC}" +from pathlib import Path +import sys + +data = Path(sys.argv[1]).read_bytes() +if b"\x39\x10" in data: + raise SystemExit("extended tmov fp form reused legacy opcode 0x1039") +if b"\xff\xff" not in data: + raise SystemExit("extended tmov fp form did not use generic v0 encoding") +PY + +"${PTOBC_BIN}" decode "${TMOV_BC}" -o "${TMOV_CURRENT_IR}" +grep -F "pto.tmov ins(" "${TMOV_CURRENT_IR}" >/dev/null +grep -F "accToVecMode = #pto.acc_to_vec_mode" "${TMOV_CURRENT_IR}" >/dev/null +grep -F "reluPreMode = #pto" "${TMOV_CURRENT_IR}" >/dev/null +"${PTOAS_BIN}" --pto-arch=a5 "${TMOV_CURRENT_IR}" -o "${TMOV_CURRENT_CPP}" +grep -F "TMOV<" "${TMOV_CURRENT_CPP}" \ + | grep -F "AccToVecMode::SingleModeVec0" \ + | grep -F "ReluPreMode::NormalRelu" >/dev/null + +TSTORE_EXTENDED_IN="${TESTDATA_DIR}/tstore_fp_extended_v0_reject.pto" +TSTORE_ERROR="${OUT_DIR}/tstore_fp_extended_v0.stderr" +if "${PTOBC_BIN}" encode "${TSTORE_EXTENDED_IN}" \ + -o "${OUT_DIR}/tstore_fp_extended_v0.ptobc" 2>"${TSTORE_ERROR}"; then + echo "error: extended tstore fp form unexpectedly encoded as PTO-BC v0" >&2 + exit 1 +fi +grep -F "cannot be represented safely in PTO-BC v0" "${TSTORE_ERROR}" >/dev/null +grep -F "legacy opcode 0x1066 would silently drop those semantics" "${TSTORE_ERROR}" >/dev/null + +if [[ -z "${LEGACY_PTOBC_BIN}" ]]; then + exit 0 +fi + +# Decode and lower with the last pre-unification PTOAS reader. This proves the +# generic tmov record retains mode/relu semantics across the version boundary. +TMOV_LEGACY_IR="${OUT_DIR}/tmov_fp_extended_v0.legacy.pto" +TMOV_LEGACY_CPP="${OUT_DIR}/tmov_fp_extended_v0.legacy.cpp" +"${LEGACY_PTOBC_BIN}" decode "${TMOV_BC}" -o "${TMOV_LEGACY_IR}" +grep -F "pto.tmov ins(" "${TMOV_LEGACY_IR}" >/dev/null +grep -F "accToVecMode = #pto.acc_to_vec_mode" "${TMOV_LEGACY_IR}" >/dev/null +grep -F "reluPreMode = #pto" "${TMOV_LEGACY_IR}" >/dev/null +"${LEGACY_PTOAS_BIN}" --pto-arch=a5 "${TMOV_LEGACY_IR}" -o "${TMOV_LEGACY_CPP}" +grep -F "TMOV<" "${TMOV_LEGACY_CPP}" \ + | grep -F "AccToVecMode::SingleModeVec0" \ + | grep -F "ReluPreMode::NormalRelu" >/dev/null + +# The simple form remains on 0x1066 and must still lower through the removed +# legacy pto.tstore_fp operation. +TSTORE_SIMPLE_BC="${OUT_DIR}/tstore_fp_simple_v0.ptobc" +TSTORE_LEGACY_IR="${OUT_DIR}/tstore_fp_simple_v0.legacy.pto" +TSTORE_LEGACY_CPP="${OUT_DIR}/tstore_fp_simple_v0.legacy.cpp" +"${PTOBC_BIN}" encode "${TESTDATA_DIR}/tstore_fp_v0_roundtrip.pto" \ + -o "${TSTORE_SIMPLE_BC}" +"${LEGACY_PTOBC_BIN}" decode "${TSTORE_SIMPLE_BC}" -o "${TSTORE_LEGACY_IR}" +grep -F "pto.tstore_fp" "${TSTORE_LEGACY_IR}" >/dev/null +"${LEGACY_PTOAS_BIN}" --pto-arch=a3 "${TSTORE_LEGACY_IR}" -o "${TSTORE_LEGACY_CPP}" +grep -F "TSTORE_FP" "${TSTORE_LEGACY_CPP}" >/dev/null diff --git a/tools/ptobc/tests/stage9_e2e.sh b/tools/ptobc/tests/stage9_e2e.sh index 65ef81ca69..74e436fb0d 100755 --- a/tools/ptobc/tests/stage9_e2e.sh +++ b/tools/ptobc/tests/stage9_e2e.sh @@ -27,6 +27,7 @@ mkdir -p "${OUT_DIR}" should_skip_roundtrip() { local path="$1" case "$path" in + */tstore_fp_extended_v0_reject.pto) return 0 ;; */test/samples/Qwen3DecodeA5/*.pto) return 0 ;; */test/samples/Complex/mix_kernel.pto) return 0 ;; */test/samples/SCF/scf_for_break_like.pto) return 0 ;; From 596559c6c007df05fe883b82f13a33a2ef6e9c34 Mon Sep 17 00:00:00 2001 From: FangRui Date: Mon, 10 Aug 2026 16:26:12 +0800 Subject: [PATCH 083/122] Fix implicit tmp for dynamic tsort32 and wide-dst tmrgsort Dynamic tsort32 valid widths were treated as 32-aligned and skipped tmp materialization, leaving the tail path without scratch at runtime; now materialize a full-width tmp at level1/2 and reject at level3. tmrgsort sized its implicit tmp from the source column sum only, which could be narrower than the destination and fail the tmp.cols >= dst.cols verifier; size it by max(totalSrcCols, dstCols) instead. --- .../Transforms/PTOMaterializeImplicitTmp.cpp | 37 ++++++++++++++++--- .../pto/tmrgsort_dst_wider_implicit_tmp.pto | 27 ++++++++++++++ .../tsort32_dynamic_width_implicit_tmp.pto | 34 +++++++++++++++++ 3 files changed, 92 insertions(+), 6 deletions(-) create mode 100644 test/lit/pto/tmrgsort_dst_wider_implicit_tmp.pto create mode 100644 test/lit/pto/tsort32_dynamic_width_implicit_tmp.pto diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp index 92139a09e6..f04e663aed 100644 --- a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -393,14 +393,30 @@ static LogicalResult replaceTSort32WithTmp(pto::TSort32Op op, if (op.getTmp()) return success(); auto valid = getValidShapeVec(op.getSrc().getType()); - if (valid.size() != 2 || valid[1] == ShapedType::kDynamic || - valid[1] % 32 == 0) + if (valid.size() != 2) + return success(); + // Only a statically-known 32-aligned width provably skips the tail path and + // needs no tmp. A dynamic width may be non-aligned at runtime, so it must be + // treated conservatively rather than assumed aligned. + bool dynamicWidth = valid[1] == ShapedType::kDynamic; + if (!dynamicWidth && valid[1] % 32 == 0) return success(); if (requireExplicitTmp) return op.emitOpError( - "requires explicit tmp for non-32-aligned tsort32 when PlanMemory is skipped"); - - FailureOr tmpType = makeSameShapeTmpType(ctx, op.getSrc()); + "requires explicit tmp for tsort32 with dynamic or non-32-aligned width " + "when PlanMemory is skipped"); + + FailureOr tmpType = failure(); + if (dynamicWidth) { + // The runtime width may be non-32-aligned, so size the scratch buffer by + // the full physical width to guarantee room for the tail path. + auto srcTy = dyn_cast(op.getSrc().getType()); + auto shape = getShapeVec(op.getSrc().getType()); + if (srcTy && shape.size() == 2 && !hasDynamicDim(shape)) + tmpType = makeVecTmpType(ctx, shape, srcTy.getElementType(), shape); + } else { + tmpType = makeSameShapeTmpType(ctx, op.getSrc()); + } if (failed(tmpType)) return op.emitOpError( "requires static tile_buf src to materialize implicit tsort32 tmp"); @@ -747,8 +763,17 @@ static LogicalResult materializeTMrgSortTmp(pto::TMrgSortOp op, } if (!elementType || totalCols <= 0) return op.emitOpError("failed to infer tmrgsort format2 tmp type"); + // The verifier requires tmp.cols >= dst.cols as well as tmp.cols >= + // sum(src.cols), so size the scratch by the wider of the two. + int64_t dstCols = 0; + for (Value dst : op.getDsts()) { + auto dstShape = getShapeVec(dst.getType()); + if (dstShape.size() == 2 && dstShape[1] != ShapedType::kDynamic) + dstCols = std::max(dstCols, dstShape[1]); + } + int64_t tmpCols = std::max(totalCols, dstCols); pto::TileBufType tmpType = - makeVecTmpType(ctx, {1, totalCols}, elementType, {1, totalCols}); + makeVecTmpType(ctx, {1, tmpCols}, elementType, {1, tmpCols}); OpBuilder builder(op); FailureOr tmp = createAllocTmp(builder, op.getLoc(), tmpType); if (failed(tmp)) diff --git a/test/lit/pto/tmrgsort_dst_wider_implicit_tmp.pto b/test/lit/pto/tmrgsort_dst_wider_implicit_tmp.pto new file mode 100644 index 0000000000..4c7739678b --- /dev/null +++ b/test/lit/pto/tmrgsort_dst_wider_implicit_tmp.pto @@ -0,0 +1,27 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// When the destination is physically wider than the sum of the source widths, +// the materialized tmp must be sized by max(sum(src.cols), dst.cols) so it +// satisfies the verifier's tmp.cols >= dst.cols constraint. Sizing it purely +// from the source column sum (128) would fail against the 192-wide dst. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @tmrgsort_dst_wider_than_srcs(%ex: vector<4xi16>) { + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tmrgsort ins(%src0, %src1 no_tmp {exhausted = false} : !pto.tile_buf, !pto.tile_buf) outs(%dst, %ex : !pto.tile_buf, vector<4xi16>) + return + } +} + +// CHECK-LABEL: func.func @tmrgsort_dst_wider_than_srcs +// CHECK: pto.tmrgsort ins(%{{.*}}, %{{.*}}, %{{.*}} {exhausted = false} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) diff --git a/test/lit/pto/tsort32_dynamic_width_implicit_tmp.pto b/test/lit/pto/tsort32_dynamic_width_implicit_tmp.pto new file mode 100644 index 0000000000..193fac9d88 --- /dev/null +++ b/test/lit/pto/tsort32_dynamic_width_implicit_tmp.pto @@ -0,0 +1,34 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// A dynamic valid width cannot be proven 32-aligned at compile time, so the +// tail path may need scratch. At level1/2 a tmp must be materialized (sized by +// the full physical width); at level3 it must be rejected rather than silently +// emitted in the 3-argument no-tmp form. + +// RUN: sed -E 's/ addr = %[A-Za-z0-9_]+//g' %s > %t.l2.pto && ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %t.l2.pto 2>&1 | FileCheck %s --check-prefix=L2 +// RUN: not ptoas --pto-arch=a3 --pto-level=level3 %s 2>&1 | FileCheck %s --check-prefix=L3 + +module { + func.func @tsort32_dynamic_width(%vr: index, %vc: index) { + %a0 = arith.constant 0 : i64 + %a1 = arith.constant 256 : i64 + %a2 = arith.constant 512 : i64 + %src = pto.alloc_tile addr = %a0 valid_row = %vr valid_col = %vc : !pto.tile_buf + %idx = pto.alloc_tile addr = %a1 valid_row = %vr valid_col = %vc : !pto.tile_buf + %dst = pto.alloc_tile addr = %a2 valid_row = %vr valid_col = %vc : !pto.tile_buf + pto.tsort32 ins(%src, %idx : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// L2-LABEL: func.func @tsort32_dynamic_width +// L2: pto.tsort32 ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + +// L3: error: 'pto.tsort32' op requires explicit tmp for tsort32 with dynamic or non-32-aligned width when PlanMemory is skipped From b3241bf461a3ac3345cc777d60a4c3767c65ade1 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Mon, 10 Aug 2026 17:06:39 +0800 Subject: [PATCH 084/122] fix: preserve FP results in PTO-BC v0 --- tools/ptobc/src/mlir_encode.cpp | 19 ++++++---- .../testdata/tmov_fp_result_v0_roundtrip.pto | 20 +++++++++++ .../testdata/tstore_fp_result_v0_reject.pto | 24 +++++++++++++ .../tests/fp_extended_v0_compatibility.sh | 35 +++++++++++++++++++ tools/ptobc/tests/stage9_e2e.sh | 1 + 5 files changed, 93 insertions(+), 6 deletions(-) create mode 100644 tools/ptobc/testdata/tmov_fp_result_v0_roundtrip.pto create mode 100644 tools/ptobc/testdata/tstore_fp_result_v0_reject.pto diff --git a/tools/ptobc/src/mlir_encode.cpp b/tools/ptobc/src/mlir_encode.cpp index d96dc129af..3ab90c6420 100644 --- a/tools/ptobc/src/mlir_encode.cpp +++ b/tools/ptobc/src/mlir_encode.cpp @@ -75,12 +75,13 @@ using FunctionVector = llvm::SmallVectorgetNumResults() == 0 && !op.getAccToVecModeAttr() && op.getReluPreMode() == mlir::pto::ReluPreMode::NoRelu; } static bool canUseLegacyTStoreFpWireOpcode(mlir::pto::TStoreOp op) { - return op.getAtomicType() == mlir::pto::AtomicType::AtomicNone && + return op->getNumResults() == 0 && + op.getAtomicType() == mlir::pto::AtomicType::AtomicNone && op.getReluPreMode() == mlir::pto::ReluPreMode::NoRelu; } @@ -101,7 +102,7 @@ static bool shouldEncodeViaGenericV0CompatibilityShim(mlir::Operation &op) { if (auto tinsert = llvm::dyn_cast(&op)) return static_cast(tinsert.getPreQuantScalar()); if (auto tmov = llvm::dyn_cast(&op)) { - if (tmov.getPreQuantScalar()) + if (tmov.getPreQuantScalar() || tmov->getNumResults() != 0) return true; // The removed pto.tmov.fp op carried only src/fp/dst. Its legacy opcode // would silently discard mode/relu semantics when read by an older PTOAS. @@ -109,6 +110,12 @@ static bool shouldEncodeViaGenericV0CompatibilityShim(mlir::Operation &op) { // so use the generic v0 record for the extended unified form. return tmov.getFp() && !canUseLegacyTMovFpWireOpcode(tmov); } + // The fixed pto.tstore schema also predates its optional tensor result. + // Result-bearing non-fp forms are readable by older unified pto.tstore + // readers through generic v0. The fp form is rejected separately because + // no pre-unification operation represented both fp and a result. + if (auto tstore = llvm::dyn_cast(&op)) + return tstore->getNumResults() != 0; if (llvm::isa( &op)) return true; @@ -150,9 +157,9 @@ getUnsupportedV0EncodingReason(mlir::Operation &op) { if (!tstore || !tstore.getFp() || canUseLegacyTStoreFpWireOpcode(tstore)) return std::nullopt; - return "pto.tstore fp with non-default atomicType or reluPreMode cannot " - "be represented safely in PTO-BC v0; legacy opcode 0x1066 would " - "silently drop those semantics"; + return "pto.tstore fp with a result or non-default atomicType/reluPreMode " + "cannot be represented safely in PTO-BC v0; legacy opcode 0x1066 " + "would silently drop those semantics"; } static bool omitsDerivedOperandSegmentsInV0(uint16_t opcode) { diff --git a/tools/ptobc/testdata/tmov_fp_result_v0_roundtrip.pto b/tools/ptobc/testdata/tmov_fp_result_v0_roundtrip.pto new file mode 100644 index 0000000000..1712755111 --- /dev/null +++ b/tools/ptobc/testdata/tmov_fp_result_v0_roundtrip.pto @@ -0,0 +1,20 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +module { + func.func @tmov_fp_result_v0() -> tensor<32x32xi8> { + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + %result = pto.tmov ins(%src : !pto.tile_buf, + %fp : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + -> tensor<32x32xi8> + return %result : tensor<32x32xi8> + } +} diff --git a/tools/ptobc/testdata/tstore_fp_result_v0_reject.pto b/tools/ptobc/testdata/tstore_fp_result_v0_reject.pto new file mode 100644 index 0000000000..913d8b0fe7 --- /dev/null +++ b/tools/ptobc/testdata/tstore_fp_result_v0_reject.pto @@ -0,0 +1,24 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +module { + func.func @tstore_fp_result_v0_reject(%dst: !pto.ptr) -> tensor<32x32xi8> { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 32 : index + %dst_tv = pto.make_tensor_view %dst, shape = [%c32, %c32], strides = [%c32, %c1] : !pto.tensor_view<32x32xi8> + %dst_part = pto.partition_view %dst_tv, offsets = [%c0, %c0], sizes = [%c32, %c32] : !pto.tensor_view<32x32xi8> -> !pto.partition_tensor_view<32x32xi8> + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %result = pto.tstore ins(%src : !pto.tile_buf + fp %fp : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<32x32xi8>) + -> tensor<32x32xi8> + return %result : tensor<32x32xi8> + } +} diff --git a/tools/ptobc/tests/fp_extended_v0_compatibility.sh b/tools/ptobc/tests/fp_extended_v0_compatibility.sh index 108a719d83..5c11a7c50d 100755 --- a/tools/ptobc/tests/fp_extended_v0_compatibility.sh +++ b/tools/ptobc/tests/fp_extended_v0_compatibility.sh @@ -54,6 +54,24 @@ grep -F "TMOV<" "${TMOV_CURRENT_CPP}" \ | grep -F "AccToVecMode::SingleModeVec0" \ | grep -F "ReluPreMode::NormalRelu" >/dev/null +TMOV_RESULT_IN="${TESTDATA_DIR}/tmov_fp_result_v0_roundtrip.pto" +TMOV_RESULT_BC="${OUT_DIR}/tmov_fp_result_v0_roundtrip.ptobc" +TMOV_RESULT_CURRENT_IR="${OUT_DIR}/tmov_fp_result_v0.current.pto" +"${PTOBC_BIN}" encode "${TMOV_RESULT_IN}" -o "${TMOV_RESULT_BC}" +"${PYTHON_EXECUTABLE}" - <<'PY' "${TMOV_RESULT_BC}" +from pathlib import Path +import sys + +data = Path(sys.argv[1]).read_bytes() +if b"\x39\x10" in data: + raise SystemExit("result-bearing tmov fp form reused resultless opcode 0x1039") +if b"\xff\xff" not in data: + raise SystemExit("result-bearing tmov fp form did not use generic v0 encoding") +PY +"${PTOBC_BIN}" decode "${TMOV_RESULT_BC}" -o "${TMOV_RESULT_CURRENT_IR}" +grep -E "%[0-9]+ = pto\.tmov " "${TMOV_RESULT_CURRENT_IR}" >/dev/null +grep -E "return %[0-9]+ : tensor<32x32xi8>" "${TMOV_RESULT_CURRENT_IR}" >/dev/null + TSTORE_EXTENDED_IN="${TESTDATA_DIR}/tstore_fp_extended_v0_reject.pto" TSTORE_ERROR="${OUT_DIR}/tstore_fp_extended_v0.stderr" if "${PTOBC_BIN}" encode "${TSTORE_EXTENDED_IN}" \ @@ -64,6 +82,16 @@ fi grep -F "cannot be represented safely in PTO-BC v0" "${TSTORE_ERROR}" >/dev/null grep -F "legacy opcode 0x1066 would silently drop those semantics" "${TSTORE_ERROR}" >/dev/null +TSTORE_RESULT_IN="${TESTDATA_DIR}/tstore_fp_result_v0_reject.pto" +TSTORE_RESULT_ERROR="${OUT_DIR}/tstore_fp_result_v0.stderr" +if "${PTOBC_BIN}" encode "${TSTORE_RESULT_IN}" \ + -o "${OUT_DIR}/tstore_fp_result_v0.ptobc" 2>"${TSTORE_RESULT_ERROR}"; then + echo "error: result-bearing tstore fp unexpectedly encoded as PTO-BC v0" >&2 + exit 1 +fi +grep -F "pto.tstore fp with a result" "${TSTORE_RESULT_ERROR}" >/dev/null +grep -F "cannot be represented safely in PTO-BC v0" "${TSTORE_RESULT_ERROR}" >/dev/null + if [[ -z "${LEGACY_PTOBC_BIN}" ]]; then exit 0 fi @@ -81,6 +109,13 @@ grep -F "TMOV<" "${TMOV_LEGACY_CPP}" \ | grep -F "AccToVecMode::SingleModeVec0" \ | grep -F "ReluPreMode::NormalRelu" >/dev/null +# The old generic reader must also reconstruct the result before decoding the +# following func.return operand. A missing result makes that value ID invalid. +TMOV_RESULT_LEGACY_IR="${OUT_DIR}/tmov_fp_result_v0.legacy.pto" +"${LEGACY_PTOBC_BIN}" decode "${TMOV_RESULT_BC}" -o "${TMOV_RESULT_LEGACY_IR}" +grep -E "%[0-9]+ = pto\.tmov " "${TMOV_RESULT_LEGACY_IR}" >/dev/null +grep -E "return %[0-9]+ : tensor<32x32xi8>" "${TMOV_RESULT_LEGACY_IR}" >/dev/null + # The simple form remains on 0x1066 and must still lower through the removed # legacy pto.tstore_fp operation. TSTORE_SIMPLE_BC="${OUT_DIR}/tstore_fp_simple_v0.ptobc" diff --git a/tools/ptobc/tests/stage9_e2e.sh b/tools/ptobc/tests/stage9_e2e.sh index 74e436fb0d..38bba806ae 100755 --- a/tools/ptobc/tests/stage9_e2e.sh +++ b/tools/ptobc/tests/stage9_e2e.sh @@ -28,6 +28,7 @@ should_skip_roundtrip() { local path="$1" case "$path" in */tstore_fp_extended_v0_reject.pto) return 0 ;; + */tstore_fp_result_v0_reject.pto) return 0 ;; */test/samples/Qwen3DecodeA5/*.pto) return 0 ;; */test/samples/Complex/mix_kernel.pto) return 0 ;; */test/samples/SCF/scf_for_break_like.pto) return 0 ;; From d90e9fc8b9bfd1c3420cbcb617b6b3194e7ae1a4 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Mon, 10 Aug 2026 18:05:11 +0800 Subject: [PATCH 085/122] fix: reject unlowerable TMOV FP results --- tools/ptobc/src/mlir_encode.cpp | 8 ++++- ...dtrip.pto => tmov_fp_result_v0_reject.pto} | 2 +- .../tests/fp_extended_v0_compatibility.sh | 34 ++++++------------- tools/ptobc/tests/stage9_e2e.sh | 1 + 4 files changed, 19 insertions(+), 26 deletions(-) rename tools/ptobc/testdata/{tmov_fp_result_v0_roundtrip.pto => tmov_fp_result_v0_reject.pto} (96%) diff --git a/tools/ptobc/src/mlir_encode.cpp b/tools/ptobc/src/mlir_encode.cpp index 3ab90c6420..7d72c84851 100644 --- a/tools/ptobc/src/mlir_encode.cpp +++ b/tools/ptobc/src/mlir_encode.cpp @@ -102,7 +102,7 @@ static bool shouldEncodeViaGenericV0CompatibilityShim(mlir::Operation &op) { if (auto tinsert = llvm::dyn_cast(&op)) return static_cast(tinsert.getPreQuantScalar()); if (auto tmov = llvm::dyn_cast(&op)) { - if (tmov.getPreQuantScalar() || tmov->getNumResults() != 0) + if (tmov.getPreQuantScalar()) return true; // The removed pto.tmov.fp op carried only src/fp/dst. Its legacy opcode // would silently discard mode/relu semantics when read by an older PTOAS. @@ -153,6 +153,12 @@ static std::optional getLegacyFpWireOpcode(mlir::Operation &op) { static std::optional getUnsupportedV0EncodingReason(mlir::Operation &op) { + if (auto tmov = llvm::dyn_cast(&op); + tmov && tmov.getFp() && tmov->getNumResults() != 0) + return "pto.tmov fp with a result cannot be represented safely in " + "PTO-BC v0; legacy opcode 0x1039 would silently drop the result, " + "and PTOAS backends do not lower the generic result-bearing form"; + auto tstore = llvm::dyn_cast(&op); if (!tstore || !tstore.getFp() || canUseLegacyTStoreFpWireOpcode(tstore)) return std::nullopt; diff --git a/tools/ptobc/testdata/tmov_fp_result_v0_roundtrip.pto b/tools/ptobc/testdata/tmov_fp_result_v0_reject.pto similarity index 96% rename from tools/ptobc/testdata/tmov_fp_result_v0_roundtrip.pto rename to tools/ptobc/testdata/tmov_fp_result_v0_reject.pto index 1712755111..c2c5175f34 100644 --- a/tools/ptobc/testdata/tmov_fp_result_v0_roundtrip.pto +++ b/tools/ptobc/testdata/tmov_fp_result_v0_reject.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. module { - func.func @tmov_fp_result_v0() -> tensor<32x32xi8> { + func.func @tmov_fp_result_v0_reject() -> tensor<32x32xi8> { %src = pto.alloc_tile : !pto.tile_buf %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf diff --git a/tools/ptobc/tests/fp_extended_v0_compatibility.sh b/tools/ptobc/tests/fp_extended_v0_compatibility.sh index 5c11a7c50d..31e79dcfbb 100755 --- a/tools/ptobc/tests/fp_extended_v0_compatibility.sh +++ b/tools/ptobc/tests/fp_extended_v0_compatibility.sh @@ -54,23 +54,16 @@ grep -F "TMOV<" "${TMOV_CURRENT_CPP}" \ | grep -F "AccToVecMode::SingleModeVec0" \ | grep -F "ReluPreMode::NormalRelu" >/dev/null -TMOV_RESULT_IN="${TESTDATA_DIR}/tmov_fp_result_v0_roundtrip.pto" -TMOV_RESULT_BC="${OUT_DIR}/tmov_fp_result_v0_roundtrip.ptobc" -TMOV_RESULT_CURRENT_IR="${OUT_DIR}/tmov_fp_result_v0.current.pto" -"${PTOBC_BIN}" encode "${TMOV_RESULT_IN}" -o "${TMOV_RESULT_BC}" -"${PYTHON_EXECUTABLE}" - <<'PY' "${TMOV_RESULT_BC}" -from pathlib import Path -import sys - -data = Path(sys.argv[1]).read_bytes() -if b"\x39\x10" in data: - raise SystemExit("result-bearing tmov fp form reused resultless opcode 0x1039") -if b"\xff\xff" not in data: - raise SystemExit("result-bearing tmov fp form did not use generic v0 encoding") -PY -"${PTOBC_BIN}" decode "${TMOV_RESULT_BC}" -o "${TMOV_RESULT_CURRENT_IR}" -grep -E "%[0-9]+ = pto\.tmov " "${TMOV_RESULT_CURRENT_IR}" >/dev/null -grep -E "return %[0-9]+ : tensor<32x32xi8>" "${TMOV_RESULT_CURRENT_IR}" >/dev/null +TMOV_RESULT_IN="${TESTDATA_DIR}/tmov_fp_result_v0_reject.pto" +TMOV_RESULT_ERROR="${OUT_DIR}/tmov_fp_result_v0.stderr" +if "${PTOBC_BIN}" encode "${TMOV_RESULT_IN}" \ + -o "${OUT_DIR}/tmov_fp_result_v0.ptobc" 2>"${TMOV_RESULT_ERROR}"; then + echo "error: result-bearing tmov fp unexpectedly encoded as PTO-BC v0" >&2 + exit 1 +fi +grep -F "pto.tmov fp with a result" "${TMOV_RESULT_ERROR}" >/dev/null +grep -F "PTOAS backends do not lower the generic result-bearing form" \ + "${TMOV_RESULT_ERROR}" >/dev/null TSTORE_EXTENDED_IN="${TESTDATA_DIR}/tstore_fp_extended_v0_reject.pto" TSTORE_ERROR="${OUT_DIR}/tstore_fp_extended_v0.stderr" @@ -109,13 +102,6 @@ grep -F "TMOV<" "${TMOV_LEGACY_CPP}" \ | grep -F "AccToVecMode::SingleModeVec0" \ | grep -F "ReluPreMode::NormalRelu" >/dev/null -# The old generic reader must also reconstruct the result before decoding the -# following func.return operand. A missing result makes that value ID invalid. -TMOV_RESULT_LEGACY_IR="${OUT_DIR}/tmov_fp_result_v0.legacy.pto" -"${LEGACY_PTOBC_BIN}" decode "${TMOV_RESULT_BC}" -o "${TMOV_RESULT_LEGACY_IR}" -grep -E "%[0-9]+ = pto\.tmov " "${TMOV_RESULT_LEGACY_IR}" >/dev/null -grep -E "return %[0-9]+ : tensor<32x32xi8>" "${TMOV_RESULT_LEGACY_IR}" >/dev/null - # The simple form remains on 0x1066 and must still lower through the removed # legacy pto.tstore_fp operation. TSTORE_SIMPLE_BC="${OUT_DIR}/tstore_fp_simple_v0.ptobc" diff --git a/tools/ptobc/tests/stage9_e2e.sh b/tools/ptobc/tests/stage9_e2e.sh index 38bba806ae..6d6b4e6e10 100755 --- a/tools/ptobc/tests/stage9_e2e.sh +++ b/tools/ptobc/tests/stage9_e2e.sh @@ -27,6 +27,7 @@ mkdir -p "${OUT_DIR}" should_skip_roundtrip() { local path="$1" case "$path" in + */tmov_fp_result_v0_reject.pto) return 0 ;; */tstore_fp_extended_v0_reject.pto) return 0 ;; */tstore_fp_result_v0_reject.pto) return 0 ;; */test/samples/Qwen3DecodeA5/*.pto) return 0 ;; From f81743307938d647d858b560216b1ab80a5a7304 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Tue, 11 Aug 2026 02:26:07 +0000 Subject: [PATCH 086/122] chore(release): bump base version to v0.58 --- CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 13ea3accf0..30e9f8ae00 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -35,7 +35,7 @@ if((NOT DEFINED MLIR_DIR OR MLIR_DIR STREQUAL "") "MLIR package directory derived from LLVM_BUILD_DIR") endif() -project(ptoas VERSION 0.57) +project(ptoas VERSION 0.58) # Wheel builds carry the same Linux hardening policy as release CMake builds. # Including the cache settings here keeps the policy backend-independent. From 7079065d4c42a3202a8d24a6cc34c6e0d30b3c31 Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 11 Aug 2026 11:19:00 +0800 Subject: [PATCH 087/122] Enlarge tci tmp in ptobc roundtrip testdata The A2/A3 tci verifier now requires the tmp capacity to be at least 1792 bytes for an i16 dst, but this ptobc roundtrip testdata still used a 512-byte tmp (1x128 f32). ptobc encode parses and verifies, so the undersized tmp made encode fail and broke both ptobc_stage9_e2e and ptobc_tci_trowexpandadd_tmp_v0_encode. Size the tci tmp to 1x512 f32 (2048 bytes) to satisfy the tightened contract. --- tools/ptobc/testdata/tci_trowexpandadd_tmp_v0_roundtrip.pto | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/ptobc/testdata/tci_trowexpandadd_tmp_v0_roundtrip.pto b/tools/ptobc/testdata/tci_trowexpandadd_tmp_v0_roundtrip.pto index 2543a45ca4..5cb8eab6f0 100644 --- a/tools/ptobc/testdata/tci_trowexpandadd_tmp_v0_roundtrip.pto +++ b/tools/ptobc/testdata/tci_trowexpandadd_tmp_v0_roundtrip.pto @@ -10,9 +10,9 @@ module { func.func @tci_trowexpandadd_tmp_v0() { %c0_i16 = arith.constant 0 : i16 - %tci_tmp = pto.alloc_tile : !pto.tile_buf + %tci_tmp = pto.alloc_tile : !pto.tile_buf %tci_dst = pto.alloc_tile : !pto.tile_buf - pto.tci ins(%c0_i16, %tci_tmp : i16, !pto.tile_buf) + pto.tci ins(%c0_i16, %tci_tmp : i16, !pto.tile_buf) outs(%tci_dst : !pto.tile_buf) %src0 = pto.alloc_tile : !pto.tile_buf From 3d798563b91e22f43a242fa2d1bfecb0f72a2a1f Mon Sep 17 00:00:00 2001 From: Melika Norouzbeygi <52199395+Melikano@users.noreply.github.com> Date: Tue, 11 Aug 2026 05:00:48 -0400 Subject: [PATCH 088/122] [PTODSL TileLib] add elemntwise 1d 2d versions (#1098) * feat(ptodsl): add elementwise 1D legality infrastructure - add a shared conservative 1D legality constraint for elementwise tiles - preserve fractal and compact-mode metadata through TileLib specialization - reject unsupported layouts, stride gaps, and mismatched logical ranges - add focused legality and daemon metadata tests - document the shared selection rules and PTO-ISA reference * fix(ptodsl): preserve deterministic template candidate ordering * feat(ptodsl): add shared element-wise 1D/2D traversal forms * feat(ptodsl): add 1D candidates for ordinary unary TileOps * feat(ptodsl): add 1D candidates for specialized unary TileOps * feat(ptodsl): add 1D candidates for ordinary Tile-Tile ops * feat(ptodsl): add 1D candidates for specialized Tile-Tile ops * feat(ptodsl): add 1D candidates for temporary Tile-Tile ops * feat(ptodsl): add 1D candidates for ordinary Tile-Scalar ops * feat(ptodsl): add 1D candidates for specialized Tile-Scalar ops * feat(ptodsl): add 1D candidate for scalar fill * feat(ptodsl): implement predicate-aware 1D/2D selection for tcmp and tcmps * feat(ptodsl): implement predicate-select 1D/2D selection for tsel and tsels * feat(ptodsl): implemented the tcvt 1D/2D conversion family * test(tilelib): add elementwise 1D/2D acceptance matrix * test(vpto): add elementwise 1D/2D runtime equivalence coverage * test(ptodsl): allow pointer arithmetic in rowwise elementwise IR * fix(ptodsl): allow flattened predicate output for tcmps * fix(vpto): preserve fusion elision for flattened elementwise loops * test(ptodsl): align tcmps expectations with flattened selection * test(vpto): fix compare 1D/2D selection coverage * fix(ptodsl): fix tfillpad and tcvt st testcases * fix(ptodsl): fix the TileOps import failure * test(vpto): exercise elementwise equivalence through TileOp expansion * fix: avoid SmallDenseSet uninitialized warning * test: use tile load/store and auto-sync in elementwise equivalence cases * fix(tilelib): default 0 fractal size to 512 and add focused test for it * fix(tilelib): restore zero-padding semantics for tfillpad * fix(tilelib): restore tcmp flat-packed predicates behaviour * fix(tilelib): export predicate legality constraints * fix(tilelib): correct predicate capacity unit comparison * fix(tests): remove duplicate pathlib import * fix(tilelib): defensively preserve candidate priority ordering * fix(tilelib): allocate tcvt 1d candidate ids dynamically * refactor(tilelib): centralize elementwise traversal metadata * refactor(tilelib): centralize flat tile legality checks * refactor(tests): simplify elementwise equivalence kernels with factory * fix(tests): validate priority-based template candidate ordering * fix(tilelib): align elementwise tests with in-process runtime --------- Co-authored-by: Zhang Zhendong --- .../ptodsl-elementwise-1d-2d-design.md | 693 ++++++ ...todsl-tilelib-template-selection-design.md | 48 +- lib/PTO/Transforms/ExpandTileOp.cpp | 13 +- .../Transforms/InsertTemplateAttributes.cpp | 50 +- .../TileFusion/PTOFusionPredicateElision.cpp | 19 +- lib/TileOps/a5/_elementwise.py | 526 +++-- lib/TileOps/a5/_remainder.py | 143 +- lib/TileOps/a5/tabs.py | 14 +- lib/TileOps/a5/tadd.py | 14 +- lib/TileOps/a5/tadds.py | 13 +- lib/TileOps/a5/tand.py | 14 +- lib/TileOps/a5/tands.py | 14 +- lib/TileOps/a5/tcmp.py | 218 +- lib/TileOps/a5/tcmps.py | 246 ++- lib/TileOps/a5/tcvt.py | 1071 ++++++++-- lib/TileOps/a5/tdiv.py | 92 +- lib/TileOps/a5/tdivs.py | 155 +- lib/TileOps/a5/texp.py | 20 +- lib/TileOps/a5/texpand.py | 26 +- lib/TileOps/a5/tfmod.py | 9 + lib/TileOps/a5/tfmods.py | 8 + lib/TileOps/a5/tlog.py | 127 +- lib/TileOps/a5/tlrelu.py | 98 +- lib/TileOps/a5/tmax.py | 14 +- lib/TileOps/a5/tmaxs.py | 13 +- lib/TileOps/a5/tmin.py | 50 +- lib/TileOps/a5/tmins.py | 13 +- lib/TileOps/a5/tmul.py | 14 +- lib/TileOps/a5/tmuls.py | 13 +- lib/TileOps/a5/tneg.py | 28 +- lib/TileOps/a5/tnot.py | 14 +- lib/TileOps/a5/tor.py | 14 +- lib/TileOps/a5/tors.py | 14 +- lib/TileOps/a5/tprelu.py | 24 +- lib/TileOps/a5/trecip.py | 92 +- lib/TileOps/a5/trelu.py | 14 +- lib/TileOps/a5/trem.py | 9 + lib/TileOps/a5/trems.py | 9 + lib/TileOps/a5/trsqrt.py | 20 +- lib/TileOps/a5/tsel.py | 241 ++- lib/TileOps/a5/tsels.py | 303 ++- lib/TileOps/a5/tshl.py | 14 +- lib/TileOps/a5/tshls.py | 13 +- lib/TileOps/a5/tshr.py | 14 +- lib/TileOps/a5/tshrs.py | 13 +- lib/TileOps/a5/tsqrt.py | 20 +- lib/TileOps/a5/tsub.py | 14 +- lib/TileOps/a5/tsubs.py | 41 +- lib/TileOps/a5/txor.py | 17 +- lib/TileOps/a5/txors.py | 19 +- .../tilelib-template-authoring.md | 187 +- ptodsl/ptodsl/_types.py | 28 +- ptodsl/ptodsl/tilelib/__init__.py | 8 + ptodsl/ptodsl/tilelib/_selection.py | 50 +- ptodsl/ptodsl/tilelib/_template_package.py | 53 + ptodsl/ptodsl/tilelib/constraints.py | 442 +++- ptodsl/ptodsl/tilelib/metadata.py | 10 +- ptodsl/ptodsl/tilelib/registry.py | 20 +- ptodsl/ptodsl/tilelib/templates/__init__.py | 9 +- ptodsl/tests/test_ptoas_runtime.py | 4 +- ptodsl/tests/test_tilelib_catalog.py | 320 ++- ptodsl/tests/test_tilelib_constraints.py | 386 +++- ptodsl/tests/test_tilelib_elementwise.py | 1888 ++++++++++++++++- ptodsl/tests/test_tilelib_select.py | 92 +- ptodsl/tests/test_tilelib_template_package.py | 71 + .../op_fusion_backend_lifecycle_level3.pto | 13 +- .../expand_tile_op_ptodsl_compare_1d_2d.pto | 174 ++ ...expand_tile_op_ptodsl_conversion_1d_2d.pto | 91 + .../expand_tile_op_ptodsl_scalar_1d_2d.pto | 89 + ...xpand_tile_op_ptodsl_scalar_fill_1d_2d.pto | 77 + .../expand_tile_op_ptodsl_select_1d_2d.pto | 221 ++ ...ile_op_ptodsl_specialized_binary_1d_2d.pto | 187 ++ ...ile_op_ptodsl_specialized_scalar_1d_2d.pto | 187 ++ ...tile_op_ptodsl_specialized_unary_1d_2d.pto | 131 ++ .../vpto/expand_tile_op_ptodsl_tabs_1d_2d.pto | 84 + test/lit/vpto/expand_tile_op_ptodsl_tadd.pto | 57 +- ..._tile_op_ptodsl_temporary_binary_1d_2d.pto | 109 + .../vpto/expand_tile_op_tilelang_tdivs.pto | 4 +- test/lit/vpto/fold_tile_buf_intrinsics.pto | 16 +- ...rt_template_attributes_candidate_order.pto | 62 + .../cases/elementwise-1d-2d-equivalence.py | 587 +++++ 81 files changed, 9405 insertions(+), 947 deletions(-) create mode 100644 docs/designs/ptodsl-elementwise-1d-2d-design.md create mode 100644 ptodsl/ptodsl/tilelib/_template_package.py create mode 100644 ptodsl/tests/test_tilelib_template_package.py create mode 100644 test/lit/vpto/expand_tile_op_ptodsl_compare_1d_2d.pto create mode 100644 test/lit/vpto/expand_tile_op_ptodsl_conversion_1d_2d.pto create mode 100644 test/lit/vpto/expand_tile_op_ptodsl_scalar_1d_2d.pto create mode 100644 test/lit/vpto/expand_tile_op_ptodsl_scalar_fill_1d_2d.pto create mode 100644 test/lit/vpto/expand_tile_op_ptodsl_select_1d_2d.pto create mode 100644 test/lit/vpto/expand_tile_op_ptodsl_specialized_binary_1d_2d.pto create mode 100644 test/lit/vpto/expand_tile_op_ptodsl_specialized_scalar_1d_2d.pto create mode 100644 test/lit/vpto/expand_tile_op_ptodsl_specialized_unary_1d_2d.pto create mode 100644 test/lit/vpto/expand_tile_op_ptodsl_tabs_1d_2d.pto create mode 100644 test/lit/vpto/expand_tile_op_ptodsl_temporary_binary_1d_2d.pto create mode 100644 test/lit/vpto/insert_template_attributes_candidate_order.pto create mode 100644 test/vpto/cases/elementwise-1d-2d-equivalence.py diff --git a/docs/designs/ptodsl-elementwise-1d-2d-design.md b/docs/designs/ptodsl-elementwise-1d-2d-design.md new file mode 100644 index 0000000000..dbcd47dfed --- /dev/null +++ b/docs/designs/ptodsl-elementwise-1d-2d-design.md @@ -0,0 +1,693 @@ +# PTODSL Element-wise 1D/2D Template Design + +## Status + +This document is the implementation inventory and design baseline for adding +deterministic 1D/2D template selection to the A5 PTODSL TileLib element-wise +operations. + +The inventory was taken on 2026-07-30 from: + +- PTOAS commit `456ebb4b6478382650d66b3ce100671b698a1e34`. +- `cann/pto-isa` commit + `57bd62714023b18d082456da2014398957e00a81`. + +The pinned A5 PTO-ISA revision is the primary reference for instruction and +ordinary 1D/2D traversal legality. The PTOAS implementation remains the source +of truth for PTODSL integration, registered operand forms, compiler metadata, +and existing generated-helper behavior. + +The shared ordinary-element-wise legality predicate and complete tile-config +metadata plumbing are now available. Ranked candidate order is also preserved +from Python through the compact candidate attribute and `ExpandTileOp`. + +`_elementwise.py` now provides reusable flattened 1D and row-wise 2D traversal +forms for ordinary unary, Tile-Tile, Tile-Scalar, and scalar-fill families. +Its registration helpers accept an explicit traversal form, derive +`loop_depth`, and attach `require_elementwise_1d(...)` to 1D candidates. All +runtime loops use the source-backed Python `for ... in range(...)` syntax and +are lowered by PTODSL's control-flow AST rewrite. + +All unary operations now support deterministic 1D/2D selection. The ordinary +operations `tabs`, `texp`, `tneg`, `tnot`, `trelu`, `trsqrt`, and `tsqrt` +register preferred shared 1D candidates while preserving their original shared +2D candidates as fallbacks. `tlog` and `trecip` keep their algorithms local to +their operation modules and use only the shared unary traversal emitters. This +raises the current unary candidate count to 20 and the current scoped candidate +count to 92. + +The ordinary Tile-Tile operations `tadd`, `tand`, `tmax`, `tmin`, `tmul`, +`tor`, `tshl`, `tshr`, and `tsub` also register preferred shared 1D candidates +and preserve their ID-0 2D fallbacks. `tmin` now uses the same shared binary +registrar instead of duplicating its row-wise loop. This raises the current +Tile-Tile candidate count to 23 and the current scoped candidate count to 101. + +The non-temporary specialized Tile-Tile operations are also migrated. `tdiv` +keeps its default and IEEE high-precision algorithms in `tdiv.py`, while +`tfmod` keeps its dtype-sensitive computation in the remainder family module; +both use shared binary traversal emitters. This raises the current Tile-Tile +candidate count to 25 and the current scoped candidate count to 103. + +At that milestone, temporary-operand Tile-Tile operations remained a +subsequent step. + +The remaining temporary-operand Tile-Tile operations are now migrated: +`tprelu`, `trem`, and `txor` preserve their existing ID-0 2D candidates and +register preferred ID-1 1D candidates. All four tile operands (`src0`, `src1`, +`tmp`, and `dst`) participate in flattened-traversal legality, including the +ABI temporary when the current generated body does not access it. This raises +the Tile-Tile candidate count to 28 and the current scoped candidate count to +106, completing 1D/2D coverage for all 14 Tile-Tile operations. + +The ordinary Tile-Scalar operations `tadds`, `tands`, `tmaxs`, `tmins`, +`tmuls`, `tors`, `tshls`, `tshrs`, and `tsubs` now also preserve their ID-0 +2D candidates and register preferred ID-1 1D candidates. Scalar operands do +not participate in memory-contiguity checks; the shared rule is applied to +`src` and `dst`. The shift operations retain their `i16` scalar signatures, +and `tsubs` now uses the shared registrar while preserving its `vbr` plus +`vsub` computation. This raises the Tile-Scalar candidate count to 24 and the +current scoped candidate count to 115. At that milestone, specialized +Tile-Scalar operations remained a subsequent step. + +The specialized Tile-Scalar operations are now migrated as well. `tdivs` +keeps both operand orders and its precision-dependent algorithm local, with +preferred 1D IDs 2 and 3 paired with existing 2D IDs 0 and 1. `tfmods` and +`trems` use the traversal-aware scalar remainder registrar, while `tlrelu` +keeps slope coercion local and `txors` retains scalar broadcasting. The +temporary tiles of `trems` and `txors` participate in 1D legality. This raises +the Tile-Scalar candidate count to 30 and the current scoped candidate count +to 121, completing 1D/2D coverage for all 14 Tile-Scalar operations. Compare, +select, conversion, scalar fill, functional execution tests, and performance +measurements remain subsequent steps. + +Scalar fill is now migrated. `texpands` preserves its ID-0 2D fallback and +registers a preferred ID-1 1D candidate for all six existing dtype +signatures. Only the destination tile participates in flattened-traversal +legality. This raises the current scoped candidate count to 122. Predicate +compare/select, conversion, functional execution tests, and performance +measurements remain subsequent steps. + +Predicate compare is now migrated. `tcmp` and `tcmps` preserve their ID-0 +row-wise fallbacks and register preferred ID-1 flattened candidates. Their +shared `require_predicate_compare_1d(...)` rule models one predicate bit per +source element together with the complete 16-byte PK or 32-byte NORM stores +used by A5. A single logical row may flatten when its physical predicate row +has enough capacity. Multiple rows additionally require the source row width +to end on a predicate-store boundary and the destination row stride to equal +the exact packed bytes produced per row. This raises the current scoped +candidate count to 124. Select, conversion, functional execution tests, and +performance measurements remain subsequent steps. + +The A5 PTOAS verifier requires `tcmp` source and destination physical shapes +to match. Since an ordinary predicate destination consequently has a wider +row stride than its packed result, practical multi-row `tcmp` cases retain the +2D candidate; eligible single-row cases select 1D. `tcmps` permits a dense +packed destination and therefore selects 1D for block-aligned multi-row cases. +The previous f32/i32 `tcmps` candidate flattened unconditionally despite its +2D metadata. It is now a genuine row-wise fallback, preventing partial rows or +predicate row padding from being crossed speculatively. + +Predicate select is now migrated. `tsel` and `tsels` preserve their ID-0 +row-wise fallbacks and register preferred ID-1 flattened candidates. Their +shared `require_predicate_select_1d(...)` rule applies the compare packing +units in reverse: ordinary data tiles must describe one contiguous logical +range, while the mask is checked as a byte-addressed packed predicate. Mask +row capacity accounts for the nominal i8/i16/i32 container dtype used by +`tsels`. Multi-row flattening requires complete predicate blocks and an exact +packed mask row stride. This raises the current scoped candidate count to 126. +Conversion, functional execution tests, and performance measurements remain +subsequent steps. + +The A5 implementations do not access the `tsel`/`tsels` temporary tile. The +shared rule nevertheless includes it in legality by requiring complete, +supported local tile metadata. Its shape is not required to match the data +range because doing so would reject the small ABI-compatible temporary tiles +used by existing callers. The f32 row-wise fallback also now derives paired +iterations from the rounded vector-repeat count, matching A5 for valid widths +between 65 and 127 instead of treating that entire range as one 64-lane tail. + +## Goals + +- Account for every A5 PTODSL element-wise TileOp in scope. +- Record existing callable forms, dtypes, attributes, temporary operands, and + traversal forms before refactoring them. +- Separate straightforward flattening from predicate- and conversion-specific + legality. +- Define a candidate identity scheme that does not make candidate IDs carry + ranking semantics. +- Expose the metadata and specialization gaps that must be closed before a 1D + candidate can be selected safely. + +## Non-Goals + +- This work does not include reductions, row/column expansion, partial updates, + load/store, data movement, sorting, random-number generation, or matrix + operations. +- This baseline does not decide that any unresolved operation is a permanent + 1D exception. +- Existing precision-mode coverage gaps are recorded but are not automatically + expanded by the traversal refactor. + +## Inventory Summary + +At the inventory baseline, the scope contained 43 TileOps and 82 registered +PTODSL candidates: + +| Family | TileOps | Registered candidates | +|---|---:|---:| +| Unary | 9 | 10 | +| Tile-Tile | 14 | 14 | +| Tile-Scalar | 14 | 15 | +| Compare, select, conversion, scalar fill | 6 | 43 | +| Total | 43 | 82 | + +At that baseline, all 82 candidates declared `loop_depth=2`. All nine unary +operations and all 14 Tile-Tile operations have since gained explicit +`loop_depth=1` candidates. The important baseline exception is `tcmps`: its +f32/i32 branch already traverses `valid_rows * valid_cols` as a flat range even +though the containing candidate is marked as two-dimensional. That mixed +candidate must be separated into explicit 1D and 2D forms. + +Each operation is loaded from +`ptodsl/ptodsl/tilelib/templates/a5/.py`, except that `texpands` is +registered by `texpand.py`. Generated registrations delegate to +`_elementwise.py` or `_remainder.py`; bespoke registrations and bodies remain +in their operation modules. The tables below name every registered candidate, +so the module and callable contract can be recovered without depending on +Python import order. + +The following abbreviations are used in the inventory: + +| Abbreviation | Dtypes | +|---|---| +| `F2` | `f16`, `f32` | +| `I6` | `i8`, `i16`, `i32`, `ui8`, `ui16`, `ui32` | +| `N9` | `I6`, `f16`, `bf16`, `f32` | +| `NEG6` | `i8`, `i16`, `i32`, `f16`, `bf16`, `f32` | +| `RELU3` | `i32`, `f16`, `f32` | +| `FILL6` | `i8`, `i16`, `i32`, `f16`, `bf16`, `f32` | + +The provisional 1D classifications mean: + +- **Shared**: the operation can use the common same-element-range legality rule + and a standard flat vector traversal. +- **Algorithm-specific**: flattening appears possible, but the operation has + computation, operand, or mode details that require a dedicated body or an + extension to the shared rule. +- **Predicate-specific**: legality depends on the packed predicate + representation as well as the data tiles. +- **Conversion-specific**: legality depends on source/destination widths, + distribution modes, packing, or a multi-step conversion. +- **Exception**: 1D has been proven illegal and the reviewed reason must be + documented and tested. No operation is classified as a confirmed exception + at this inventory stage. + +## Unary Inventory + +| Op | Existing candidate(s) and operands | Dtypes | Current form | Provisional 1D classification | +|---|---|---|---|---| +| `tabs` | `template_tabs(src, dst)`; `template_tabs_1d(src, dst)` | same `F2` | shared 2D fallback and preferred shared 1D | Shared | +| `texp` | `template_texp(src, dst)`; `template_texp_1d(src, dst)` | same `F2` | shared 2D fallback and preferred shared 1D, default precision only | Shared | +| `tlog` | default `template_tlog`/`template_tlog_1d`; high-precision `template_tlog_high_precision`/`template_tlog_high_precision_1d` | same `F2` | precision-selected 2D fallbacks and preferred 1D candidates | Algorithm-specific | +| `tneg` | `template_tneg(src, dst)`; `template_tneg_1d(src, dst)` | same `NEG6` | shared 2D fallback and preferred shared 1D | Shared | +| `tnot` | `template_tnot(src, dst)`; `template_tnot_1d(src, dst)` | same `I6` | shared 2D fallback and preferred shared 1D | Shared | +| `trecip` | `template_trecip(src, dst)`; `template_trecip_1d(src, dst)` | same `F2` | operation-local `1 / src` computation with shared 2D/1D traversal | Algorithm-specific | +| `trelu` | `template_trelu(src, dst)`; `template_trelu_1d(src, dst)` | same `RELU3` | shared 2D fallback and preferred shared 1D | Shared | +| `trsqrt` | `template_trsqrt(src, dst)`; `template_trsqrt_1d(src, dst)` | same `F2` | shared 2D fallback and preferred shared 1D, default precision only | Shared | +| `tsqrt` | `template_tsqrt(src, dst)`; `template_tsqrt_1d(src, dst)` | same `F2` | shared 2D fallback and preferred shared 1D, default precision only | Shared | + +`precisionType` is forwarded for `texp`, `tlog`, `trecip`, `trsqrt`, and +`tsqrt`. In the current PTODSL templates only `tlog` uses it for candidate +selection; the other listed templates explicitly implement default precision +only. The 1D/2D work must preserve that current coverage unless precision +parity is approved as separate work. + +## Tile-Tile Inventory + +| Op | Existing candidate and operands | Dtypes | Current form | Provisional 1D classification | +|---|---|---|---|---| +| `tadd` | `template_tadd`/`template_tadd_1d(src0, src1, dst)` | same `N9` | shared 2D fallback and preferred shared 1D | Shared | +| `tand` | `template_tand`/`template_tand_1d(src0, src1, dst)` | same `I6` | shared 2D fallback and preferred shared 1D | Shared | +| `tdiv` | `template_tdiv`/`template_tdiv_1d(src0, src1, dst)` | same `F2` | operation-local precision-aware computation with shared 2D/1D traversal | Algorithm-specific | +| `tfmod` | `template_tfmod`/`template_tfmod_1d(src0, src1, dst)` | same `f32`, `f16`, `i16`, or `ui16` | remainder-family 2D fallback and preferred 1D | Algorithm-specific | +| `tmax` | `template_tmax`/`template_tmax_1d(src0, src1, dst)` | same `N9` | shared 2D fallback and preferred shared 1D | Shared | +| `tmin` | `template_tmin`/`template_tmin_1d(src0, src1, dst)` | same `N9` | shared 2D fallback and preferred shared 1D | Shared | +| `tmul` | `template_tmul`/`template_tmul_1d(src0, src1, dst)` | same `N9` | shared 2D fallback and preferred shared 1D | Shared | +| `tor` | `template_tor`/`template_tor_1d(src0, src1, dst)` | same `I6` | shared 2D fallback and preferred shared 1D | Shared | +| `tprelu` | `template_tprelu`/`template_tprelu_1d(src0, src1, tmp, dst)` | data/dst `f16` or `f32`; `tmp` is the data dtype or `i8` | shared binary 2D fallback and preferred 1D; temporary included in legality | Algorithm-specific | +| `trem` | `template_trem`/`template_trem_1d(src0, src1, tmp, dst)` | all `f32`, all `f16`, or all `i32` | remainder-family 2D fallback and preferred 1D; temporary included in legality | Algorithm-specific | +| `tshl` | `template_tshl`/`template_tshl_1d(src0, src1, dst)` | same `I6` | shared 2D fallback and preferred shared 1D | Shared | +| `tshr` | `template_tshr`/`template_tshr_1d(src0, src1, dst)` | same `I6` | shared 2D fallback and preferred shared 1D | Shared | +| `tsub` | `template_tsub`/`template_tsub_1d(src0, src1, dst)` | same `N9` | shared 2D fallback and preferred shared 1D | Shared | +| `txor` | `template_txor`/`template_txor_1d(src0, src1, tmp, dst)` | same `I6` across all operands | shared binary 2D fallback and preferred 1D; temporary included in legality | Algorithm-specific | + +`precisionType` is forwarded and consumed by `tdiv`. A 1D form must preserve +both its ordinary `vdiv` path and its high-precision helper path. + +Every temporary tile remains part of the 1D legality decision even where the +current PTODSL body does not read it directly. Whether a temporary's logical +range must equal the data range or satisfy a representation-specific relation +is an operation-family decision; it must not be omitted from the predicate. + +## Tile-Scalar Inventory + +| Op | Existing candidate(s) and operands | Dtypes | Current form | Provisional 1D classification | +|---|---|---|---|---| +| `tadds` | `template_tadds`/`template_tadds_1d(src, scalar, dst)` | same `N9` | shared 2D fallback and preferred shared 1D | Shared | +| `tands` | `template_tands`/`template_tands_1d(src, scalar, dst)` | same `I6` | shared 2D fallback and preferred shared 1D | Shared | +| `tdivs` | `template_tdivs_tile_scalar`/`template_tdivs_tile_scalar_1d(src, scalar, dst)`; `template_tdivs_scalar_tile`/`template_tdivs_scalar_tile_1d(scalar, src, dst)` | same `F2` | operand-order-specific precision-aware 2D fallbacks and preferred 1D candidates | Algorithm-specific | +| `tfmods` | `template_tfmods`/`template_tfmods_1d(src, scalar, dst)` | same `f32`, `f16`, `i32`, or `i16` | scalar remainder 2D fallback and preferred 1D | Algorithm-specific | +| `tlrelu` | `template_tlrelu`/`template_tlrelu_1d(src, slope, dst)` | `(f16,f16,f16)`, `(f16,f32,f16)`, `(f32,f32,f32)` | operation-local slope coercion with shared 2D/1D traversal | Algorithm-specific | +| `tmaxs` | `template_tmaxs`/`template_tmaxs_1d(src, scalar, dst)` | same `N9` | shared 2D fallback and preferred shared 1D | Shared | +| `tmins` | `template_tmins`/`template_tmins_1d(src, scalar, dst)` | same `N9` | shared 2D fallback and preferred shared 1D | Shared | +| `tmuls` | `template_tmuls`/`template_tmuls_1d(src, scalar, dst)` | same `N9` | shared 2D fallback and preferred shared 1D | Shared | +| `tors` | `template_tors`/`template_tors_1d(src, scalar, dst)` | same `I6` | shared 2D fallback and preferred shared 1D | Shared | +| `trems` | `template_trems`/`template_trems_1d(src, scalar, tmp, dst)` | all `f32` or all `f16` | scalar remainder 2D fallback and preferred 1D; temporary included in legality | Algorithm-specific | +| `tshls` | `template_tshls`/`template_tshls_1d(src, scalar, dst)` | data/dst `I6`; scalar `i16` | shared 2D fallback and preferred shared 1D | Shared | +| `tshrs` | `template_tshrs`/`template_tshrs_1d(src, scalar, dst)` | data/dst `I6`; scalar `i16` | shared 2D fallback and preferred shared 1D | Shared | +| `tsubs` | `template_tsubs`/`template_tsubs_1d(src, scalar, dst)` | same `N9` | shared broadcast-scalar 2D fallback and preferred 1D | Shared | +| `txors` | `template_txors`/`template_txors_1d(src, scalar, tmp, dst)` | same `I6` across all operands | shared scalar 2D fallback and preferred 1D; temporary included in 1D legality | Algorithm-specific | + +`tdivs` is the only scoped Tile-Scalar operation with both normal and reverse +operand forms currently registered. Both forms must receive distinct 1D +candidates and must retain their operand-kind constraints. Adding reverse forms +to other operations is outside this work unless a missing existing PTO callable +form is demonstrated separately. + +## Compare, Select, and Scalar-Fill Inventory + +| Op | Existing candidate and operands | Dtypes | Current form | Provisional 1D classification | +|---|---|---|---|---| +| `tcmp` | `template_tcmp(src0, src1, dst)` | data pairs `f32`, `i32`, `f16`, `i16`, `i8`, or `ui8`; predicate destination `i8` | dtype-specific packed predicate stores inside 2D row traversal | Predicate-specific | +| `tcmps` | `template_tcmps(src, scalar, dst)` | source/scalar `f32`, `i32`, `f16`, `i16`, `i8`, or `ui8`; predicate destination `ui8` | mixed: f32/i32 flatten, other dtypes traverse rows; metadata says 2D | Predicate-specific | +| `tsel` | `template_tsel(mask, src0, src1, tmp, dst)` | mask `i8`; data/tmp/dst all `f32`, all `f16`, or all `i8` | dtype-specific predicate loads inside 2D row traversal | Predicate-specific | +| `tsels` | `template_tsels(mask, src, tmp, scalar, dst)` | mask `i8`, `i16`, or `i32`; data/tmp/scalar/dst all one of `i8`, `i16`, `i32`, `f32`, or `f16` | dtype-specific predicate loads inside 2D row traversal | Predicate-specific | +| `texpands` | `template_texpands`/`template_texpands_1d(scalar, dst)` | same `FILL6` | shared scalar-fill 2D fallback and preferred 1D | Shared | + +The existing shape rules do not yet express predicate representation: + +- `tcmp` currently requires the predicate destination and data inputs to have + the same valid shape even though the body stores packed predicate bytes. +- `tcmps` does not require source and destination valid shapes to match. Its + f32/i32 branch flattens the source unconditionally, so a multi-row + partial-column source can cross a physical row gap. +- `tsel` checks no relationship between the mask range and the data range. +- `tsels` checks `src` and `dst` equality but deliberately ignores `mask` and + `tmp` shapes. + +These are inventory findings, not permission to tighten public legality +silently. The predicate-specific design must establish the intended physical +mask contract and add regression coverage before changing these constraints. + +`cmp_mode` is forwarded and consumed by `tcmp` and `tcmps`. It changes the +generated helper body and therefore remains part of specialization. + +## Conversion Inventory + +`tcvt` has 38 existing candidates. Every candidate has operands `(src, dst)`, +priority 0, a unique ID from 0 through 37, and `loop_depth=2`. Except for the +packed BF16-to-FP4 form, the current legality predicate requires equal physical +and valid shapes with row-major/none-box layouts. `round_mode` is forwarded and +used by conversion bodies that request rounding. + +| Op | Existing candidates and operands | Dtypes | Current form | Provisional 1D classification | +|---|---|---|---|---| +| `tcvt` | 38 candidates, each `(src, dst)` | 40 registered signatures across 38 candidates, listed below | conversion-specific 2D row traversal | Conversion-specific | + +| Existing ID | Candidate | Dtype signature | Implementation category | +|---:|---|---|---| +| 0 | `template_tcvt_f32_to_i32` | `f32 -> i32` | generic conversion | +| 1 | `template_tcvt_i32_to_f32` | `i32 -> f32` | generic conversion | +| 2 | `template_tcvt_i16_to_f16` | `i16 -> f16` | generic conversion | +| 3 | `template_tcvt_f16_to_i16` | `f16 -> i16` | two-step conversion, unpack/pack | +| 4 | `template_tcvt_bf16_to_f16` | `bf16 -> f16` | generic conversion | +| 5 | `template_tcvt_f32_to_f16` | `f32 -> f16` | packed store | +| 6 | `template_tcvt_f32_to_bf16` | `f32 -> bf16` | packed store | +| 7 | `template_tcvt_f16_to_i32` | `f16 -> i32` | unpacked load | +| 8 | `template_tcvt_f16_to_f32` | `f16 -> f32` | unpacked load | +| 9 | `template_tcvt_bf16_to_i32` | `bf16 -> i32` | unpacked load | +| 10 | `template_tcvt_ui8_to_ui16` | `ui8 -> ui16` | unpacked load | +| 11 | `template_tcvt_f32_to_fp8` | `f32 -> f8e4m3` or `f8e5m2` | low-precision select/reorder | +| 12 | `template_tcvt_f32_to_hif8` | `f32 -> hif8` | low-precision select/reorder | +| 13 | `template_tcvt_f16_to_hif8` | `f16 -> hif8` | packed low-precision store | +| 14 | `template_tcvt_bf16_to_fp4` | `bf16 -> f4e1m2x2` or `f4e2m1x2` | packed 2:1 source/destination columns | +| 15 | `template_tcvt_f32_to_i16` | `f32 -> i16` | two-step conversion, packed store | +| 16 | `template_tcvt_f32_to_i64` | `f32 -> i64` | widening store representation | +| 17 | `template_tcvt_f32_to_f32` | `f32 -> f32` | truncation operation | +| 18 | `template_tcvt_f16_to_ui8` | `f16 -> ui8` | packed store | +| 19 | `template_tcvt_f16_to_si8` | `f16 -> si8` | multi-step conversion, packed store | +| 20 | `template_tcvt_bf16_to_f32` | `bf16 -> f32` | unpacked load | +| 21 | `template_tcvt_i16_to_f32` | `i16 -> f32` | unpacked load | +| 22 | `template_tcvt_i16_to_i32` | `i16 -> i32` | unpacked load | +| 23 | `template_tcvt_i16_to_ui32` | `i16 -> ui32` | unpacked load | +| 24 | `template_tcvt_ui8_to_f16` | `ui8 -> f16` | unpacked load | +| 25 | `template_tcvt_si8_to_f16` | `si8 -> f16` | unpacked load | +| 26 | `template_tcvt_si8_to_si16` | `si8 -> si16` | unpacked load, 16-bit store mode | +| 27 | `template_tcvt_i32_to_i64` | `i32 -> i64` | widening store representation | +| 28 | `template_tcvt_i32_to_i16` | `i32 -> i16` | packed store | +| 29 | `template_tcvt_i32_to_ui16` | `i32 -> ui16` | saturated packed store | +| 30 | `template_tcvt_ui32_to_i16` | `ui32 -> i16` | saturated packed store | +| 31 | `template_tcvt_ui32_to_ui16` | `ui32 -> ui16` | saturated packed store | +| 32 | `template_tcvt_si8_to_i32` | `si8 -> i32` | interleave plus two output stores | +| 33 | `template_tcvt_i32_to_ui8` | `i32 -> ui8` | select/reorder plus byte store | +| 34 | `template_tcvt_ui32_to_ui8` | `ui32 -> ui8` | select/reorder plus byte store | +| 35 | `template_tcvt_i16_to_ui8` | `i16 -> ui8` | packed store | +| 36 | `template_tcvt_i64_to_f32` | `i64 -> f32` | packed 64-to-32 store | +| 37 | `template_tcvt_i64_to_i32` | `i64 -> i32` | packed 64-to-32 store | + +All 38 `tcvt` candidates now have a paired preferred 1D form. The A5 +implementation provides flattened helpers for the generic, unpacked-load, +packed-store, multi-step, widening, 64-bit, and low-precision categories, so +none of the currently registered PTODSL conversions is a 2D-only exception. +BF16-to-FP4 uses a separate rule that proves its 2:1 source/destination column +relationship over both physical and valid ranges. + +## Existing Shared Traversal Structure + +The ordinary implementations are concentrated in: + +- `ptodsl/ptodsl/tilelib/templates/a5/_elementwise.py` +- `ptodsl/ptodsl/tilelib/templates/a5/_remainder.py` +- `ptodsl/ptodsl/tilelib/templates/a5/_common.py` + +`_elementwise.py` registers unary, Tile-Tile, Tile-Scalar, and scalar-fill +templates. Unmigrated operation call sites retain the default 2D form, which +restarts `remained = valid_cols` for every row and then walks columns by vector +lane count. Migrated unary and Tile-Tile operations add an explicit 1D +registration next to that unchanged fallback. Specialized algorithms remain +outside `_elementwise.py` and call its lower-level family traversal emitters. +`_remainder.py` owns remainder/fmod computation and exposes traversal-aware +binary registration without moving that computation into the ordinary module. + +The reusable ordinary traversal foundation consists of: + +- `emit_elementwise_1d(anchor, emit_chunk)`, which computes + `valid_rows * valid_cols`, carries one remaining-element count, and invokes + the chunk emitter from one Python range loop; +- `emit_elementwise_2d(anchor, emit_chunk)`, which retains the outer row loop + and starts a new remaining-column count for each row; +- unary, Tile-Tile, Tile-Scalar, and scalar-fill wrappers for both forms; +- registration helpers whose `traversal="1d"` form derives `loop_depth=1`, + adds the named-operand 1D legality predicate, and whose `traversal="2d"` + form derives `loop_depth=2`. + +The 1D family wrappers use a base tile pointer plus the flattened element +offset. They do not reconstruct row and column indices. The 2D wrappers retain +row/column addressing so each physical tile's row stride remains respected. +The two module-level traversal cores opt into `rewrite_jit_function`; this is +required because only a registered template body and its lexically nested +helpers are rewritten automatically. +Candidate IDs and priorities are explicit registrar parameters, allowing a 1D +candidate to rank ahead of its stable-ID 2D fallback without relying on +registration order. + +The operation-specific files that cannot be migrated by replacing a shared +registrar alone are: + +- `tdiv.py`, `tdivs.py`, `tlog.py`, `trecip.py`, `tlrelu.py`, and `tsubs.py`; +- `tcmp.py`, `tcmps.py`, `tsel.py`, and `tsels.py`; +- `tcvt.py`. + +Temporary-operand forms in `tprelu`, `trem`, `trems`, `txor`, and `txors` also +require explicit audit even when their bodies come from a shared registrar. + +### Existing constraint patterns + +The current registrations use several related but non-identical legality +patterns: + +- Ordinary `_elementwise.py` forms accept tile operands in `ub` or `vec`, + require row-major/none-box layouts, and require the named data/temporary + operands to have equal valid shapes. +- `_remainder.py` forms require `ub`, row-major/none-box tiles and equal valid + shapes, including the temporary tile when present. +- `tdiv` constrains block layout and memory space through metadata but does not + currently require equal valid shapes or `none_box` sub-layout. +- `trecip` requires `ub`, row-major/none-box tiles but does not currently + require equal source/destination valid shapes. +- `txors` checks equal valid shapes only for `src` and `dst`; its `tmp` tile is + constrained by location/layout but omitted from the shape relation. +- `tcvt` checks physical shape, valid shape, block layout, and sub-layout in its + custom predicates, but does not declare a common memory-space restriction. + +The new shared 1D predicate must compose with each operation's existing legal +domain. It must not accidentally make the 2D fallback narrower while adding a +more conservative 1D candidate. + +## Pinned PTO-ISA Reference + +At the pinned PTO-ISA revision: + +- `include/pto/npu/a5/TUnaryOp.hpp` provides separate + `TUnaryOps_1D_*` and `TUnaryOps_2D` implementations. +- `include/pto/npu/a5/TBinOp.hpp` provides separate Tile-Tile 1D and 2D + implementations. +- `include/pto/npu/a5/TBinSOp.hpp` provides separate Tile-Scalar 1D and 2D + implementations. + +Their compile-time selection admits the 1D path when all participating ordinary +tiles have full valid columns, or when all physical tiles have one row. The +PTO-ISA types also expose `RowStride`, and its 2D forms use those row strides +explicitly. + +PTODSL follows this ordinary legality model and makes the implicit tile-type +conditions explicit: + +- every named traversal operand is checked, including temporary tiles; +- the tiles must describe the same static logical valid shape; +- row-major/none-box local tiles are gap-free when their compact mode is + `null` or `normal`; +- every tile must use its full physical column axis, or the logical valid + region must occupy only the first row; +- unknown information is rejected; +- predicate and conversion representations remain separate family rules. + +Pinned references: + +- `https://gitcode.com/cann/pto-isa/blob/57bd62714023b18d082456da2014398957e00a81/include/pto/npu/a5/TUnaryOp.hpp` +- `https://gitcode.com/cann/pto-isa/blob/57bd62714023b18d082456da2014398957e00a81/include/pto/npu/a5/TBinOp.hpp` +- `https://gitcode.com/cann/pto-isa/blob/57bd62714023b18d082456da2014398957e00a81/include/pto/npu/a5/TBinSOp.hpp` + +## Shared Ordinary 1D Legality Infrastructure + +`tilelib.require_elementwise_1d(*operand_names)` implements the reusable +ordinary-tile rule. It is intentionally opt-in: operation registrations must +name every tile that participates in traversal, including ABI temporary tiles. +The predicate accepts only when all of the following are proven: + +1. Every named operand is a rank-2 tile with positive static physical and valid + shapes. +2. Every valid extent is within its physical extent. +3. Every tile is in local `ub`/`vec` memory with `row_major` block layout, + `none_box` sub-layout, and a gap-free compact mode (`null` or `normal`). +4. All named tiles have the same logical valid shape. +5. Either every tile has `valid_cols == physical_cols`, or the common logical + valid row count is one. + +The logical one-row case is safe even when the physical tile has additional +rows because a TileBuf valid region begins at the first element and never +crosses a row boundary. Multi-row partial-column regions are rejected. + +This predicate is not sufficient for predicate tiles or dtype-width-changing +and packed conversions. Predicate compare uses the separate rule below; +predicate select (`tsel` and `tsels`) and conversion (`tcvt`) must establish +their logical-to-physical representation rules before registering a 1D +candidate. + +### Packed predicate compare legality + +`tilelib.require_predicate_compare_1d(*data_operand_names, +predicate_operand="dst")` implements the compare-specific rule. The data +tiles must satisfy the same rank, locality, layout, compact-mode, bounds, and +common-valid-range requirements as ordinary element-wise operands. The +predicate destination is checked independently because it stores one bit per +comparison instead of one element of the source dtype. + +For f32/i32 and f16/i16 comparisons, one complete predicate store represents +128 source elements and occupies 16 bytes. For i8/ui8 comparisons, one store +represents 256 source elements and occupies 32 bytes. A one-row range is legal +when every source physical row can contain the rounded full vector loads and +the destination physical row can contain the rounded number of stores. +For multiple rows, the source valid column count must be a multiple of the +corresponding 128- or 256-element store unit, every data tile must use its full +physical column axis, and the destination physical row stride must equal the +exact packed byte count for that row. The predicate physical row must also +satisfy A5's 32-byte tile-row alignment. Unknown dtype, shape, layout, compact +mode, alignment, or insufficient destination capacity rejects the 1D +candidate. + +### Packed predicate select legality + +`tilelib.require_predicate_select_1d(predicate_operand, +*data_operand_names, temporary_operand=...)` implements the inverse packed +representation rule for `tsel` and `tsels`. All data operands must have one +common static valid shape and satisfy local row-major, none-box, gap-free +continuity. The mask row byte width is `mask.shape[1] * bytewidth(mask.dtype)`; +the nominal mask dtype is a storage container and does not change the one-bit +predicate interpretation. + +The same 128-element/16-byte unit is used for 32- and 16-bit data, and the same +256-element/32-byte unit is used for 8-bit data. A single row requires enough +rounded data and mask capacity. Multiple rows additionally require full data +columns, a valid column count ending on the relevant predicate unit, and a +mask row stride equal to the exact packed bytes per row. The mask physical row +must satisfy 32-byte alignment. The A5-unused temporary is checked for static +local layout and compact-mode metadata but does not share the data range. + +The older `require_contiguous()` helper remains available for existing users; +new element-wise 1D candidates must use the named-operand rule. + +### Conversion legality + +`tilelib.require_conversion_1d()` proves that the source and destination are +two independently contiguous typed streams. Both operands must be static +rank-2 local tiles with row-major, none-box, gap-free compact storage. Ordinary +conversions require equal physical and valid shapes. Multi-row regions must +fill both physical column axes; a single logical row may use partial columns +because neither stream crosses a row boundary. + +Changing dtype width does not by itself make flattening illegal. Each source +and destination pointer advances in its own element type, while the existing +unpack/pack distribution mode and mask representation preserve the conversion +body's logical-element mapping. BF16-to-FP4 supplies +`source_elements_per_destination=2`, requiring source columns to be exactly +twice the packed destination columns. Unknown memory, layout, shape, or compact +metadata rejects the 1D candidate. + +The current A5 PTODSL candidates have only `(src, dst)` operands. A5's +non-saturating multi-step conversions use register temporaries, so no ABI +temporary tile is omitted from the proof. A future candidate with a tile +temporary must extend the conversion rule to validate that operand before it +can gain a 1D form. + +## Metadata and Legality Status + +### Current continuity predicate + +`tilelib.require_contiguous()` returns true when all tile valid +columns equal their physical columns, or when all physical row counts equal +one. It does not: + +- require equal logical valid shapes; +- inspect valid row counts for the single-row case; +- name the exact participating operands; +- represent an explicit tile row stride; +- model predicate or conversion packing; +- distinguish unknown metadata from a proven-contiguous range. + +It is therefore not the shared 1D legality rule. + +### Tile specialization metadata + +The compiler-side operand JSON includes physical shape, valid shape, memory +space, block layout, sub-layout, sub-fractal size, pad value, and compact mode. +The compiler-side `SpecKey` also includes these fields. + +Python `TileSpec`, daemon reconstruction, rendered `tile_buf` types, and the +compiler JSON path now preserve `s_fractal_size` and `compact_mode`. The +`ExpandTileOp` specialization key and generated helper name also distinguish +compact mode. + +PTO tile layout lowering establishes that an unboxed row-major tile has +`col_stride == 1`. Its row stride equals the physical column count for compact +mode `null` or `normal`; `row_plus_one` adds a stride gap. The shared predicate +therefore accepts the first two modes and rejects `row_plus_one` or unknown +compact metadata. No explicit row-stride field is needed for this restricted +ordinary layout contract. + +### Candidate ordering + +The Python registry and daemon sort legal candidates by descending priority. +Equal-priority reporting is canonicalized by candidate name, while an equal +top-priority match is rejected as ambiguous. + +The daemon metadata response represents candidates as a ranked JSON array. +`InsertTemplateAttributes` validates unique candidate IDs without reordering +the array, stores the same order in the compact candidates attribute, and +`ExpandTileOp` selects candidate zero by name. Candidate IDs are stable identity +fields only and do not participate in ranking. + +A compiler regression supplies a preferred candidate with ID 99 and a fallback +with ID 0. It verifies both the compact attribute order and that expansion +requests the ID-99 candidate, preventing an accidental return to ID-based +selection. + +### Context attributes + +The two compiler stages both forward the scoped body-changing context +attributes: + +- `round_mode` for `tcvt`; +- `cmp_mode` for `tcmp` and `tcmps`; +- `precisionType` for the precision-aware unary/division families. + +The selection and specialization work must keep the two C++ reconstructions in +lockstep. Candidate filtering may use a context attribute directly, as `tlog` +does, or the rendered body may consume it, as `tdiv`, `tdivs`, `tcmp`, `tcmps`, +and `tcvt` do. + +## Proposed Candidate Identity Scheme + +Candidate identity and ranking must remain separate. + +1. Preserve every existing candidate name and ID for its 2D fallback. This + minimizes churn in named-render tests and keeps current IDs stable. +2. Add `_1d` to each new flattened candidate name. +3. Give a legal 1D candidate a higher priority than its 2D fallback. +4. Use `loop_depth=1` for flattened candidates. +5. Retain `loop_depth=2` for fallback candidates. +6. Never rely on the numeric ID to select the winner. + +Proposed ID allocation: + +| Existing form | Preserved 2D IDs | New 1D IDs | +|---|---|---| +| Ordinary one-candidate op | `0` | `1` | +| `tlog` default/high-precision | `0`, `1` | `2`, `3`, paired in the same order | +| `tdivs` tile-scalar/scalar-tile | `0`, `1` | `2`, `3`, paired in the same order | +| `tcvt` semantic variants | `0` through `37` | `38 + existing_id`, giving `38` through `75` | + +This allocation deliberately makes the preferred 1D IDs larger than the +fallback IDs, so an accidental ascending-ID sort is exposed by tests rather +than appearing to work. + +## Resolved Conversion Decisions + +1. The registered A5 PTODSL `tcvt` forms contain only source and destination + tiles; their multi-step paths use register temporaries. +2. Width-changing forms use equal logical element offsets on independently + typed pointers. Their existing load/store distributions define the bytes + consumed and produced per iteration. +3. A5 supplies flat forms for every registered packed and multi-step category; + no form requires per-row state when both streams are proven contiguous. +4. `pad_value` does not change conversion selection because the 1D body stores + only the mask-bounded logical range and does not consume padding values as + operands. +5. Non-static shapes and unknown memory, layout, sub-layout, or compact-mode + metadata conservatively reject the 1D candidate and retain the 2D fallback. + +The implementation follows PTO-ISA revision +`23e31ddf51233835810997ba7cff12fda2808f50`, principally +`include/pto/common/arch/register/tcvt_common.hpp`, which exposes paired 1D and +2D conversion helpers and selects the flat path for full-column or single-row +tiles. + +## Coverage Checklist + +The inventory accounts for every operation in the issue exactly once: + +- Unary: `tabs`, `texp`, `tlog`, `tneg`, `tnot`, `trecip`, `trelu`, `trsqrt`, + `tsqrt`. +- Tile-Tile: `tadd`, `tand`, `tdiv`, `tfmod`, `tmax`, `tmin`, `tmul`, `tor`, + `tprelu`, `trem`, `tshl`, `tshr`, `tsub`, `txor`. +- Tile-Scalar: `tadds`, `tands`, `tdivs`, `tfmods`, `tlrelu`, `tmaxs`, + `tmins`, `tmuls`, `tors`, `trems`, `tshls`, `tshrs`, `tsubs`, `txors`. +- Other scoped element-wise families: `tcmp`, `tcmps`, `tsel`, `tsels`, + `tcvt`, `texpands`. + +Future catalog tests should encode this list as data and require each operation +to provide: + +- a legal 2D fallback; +- a preferred legal 1D candidate for eligible metadata; or +- a named, reviewed, and tested exception record. diff --git a/docs/designs/ptodsl-tilelib-template-selection-design.md b/docs/designs/ptodsl-tilelib-template-selection-design.md index f54115dd78..ec9c26def6 100644 --- a/docs/designs/ptodsl-tilelib-template-selection-design.md +++ b/docs/designs/ptodsl-tilelib-template-selection-design.md @@ -118,7 +118,7 @@ from MLIR. The JSON shape sent to the Python service is deliberately close to | Operand kind | Required metadata | |---|---| -| tile | dtype, shape, valid shape, memory space, block layout, sub-layout, fractal size, pad value | +| tile | dtype, shape, valid shape, memory space, block layout, sub-layout, fractal size, pad value, compact mode | | view | dtype, shape, strides when known, memory space, optional layout | | vector | dtype and vector shape | | scalar | dtype and static integer value when recoverable | @@ -159,21 +159,36 @@ then evaluates each registered candidate: 5. Check layout and memory-space metadata. 6. Merge context attributes. 7. Run custom constraint predicates. -8. Sort legal candidates by descending priority. +8. Sort legal candidates by descending priority, using name only to make + equal-priority reporting deterministic. If no candidate is legal, the service reports a `NoMatchingTemplate` error with per-candidate reasons. If multiple candidates tie for the highest priority and -no explicit candidate is requested, the registry reports ambiguity rather than -silently picking one. - -For multi-candidate ops, candidate `id` values must be unique. The C++ pass -sorts persisted candidate metadata by `id` and then by name, so ids should be -stable and intentionally assigned. +no explicit candidate is requested, both normal selection and metadata +insertion report ambiguity rather than silently picking one. + +Constraint predicates may depend on concrete operand metadata, so general +overlap cannot be proven when templates are merely registered. In-tree catalog +selection tests catch ties for their representative operand forms before +compiler integration runs; metadata insertion retains the concrete check for +forms that a catalog cannot exhaustively enumerate. The ambiguity diagnostic +directs authors to assign distinct priorities or make the constraints mutually +exclusive. + +For multi-candidate ops, candidate `id` values must be unique and stable. IDs +identify versions; they do not rank them. + +The service returns legal candidates as a JSON array in Python ranking order. +Each wire entry includes priority so `InsertTemplateAttributes` can defensively +normalize the result by descending priority and reject a highest-priority tie. +This protects selection across the compiler/service boundary if a response is +unsorted. Candidate ID, JSON object order, registration order, and import order +do not participate in ranking. ## Compact Candidate Attribute -`InsertTemplateAttributes` stores a compact `candidates` array attribute on the -TileOp. Each entry contains: +`InsertTemplateAttributes` stores the normalized candidates as a compact +`candidates` array attribute on the TileOp. Each entry contains: - `id` - `name` @@ -182,9 +197,11 @@ TileOp. Each entry contains: - `tail` This attribute is intentionally not a copy of the full Python metadata object. -Legality has already happened in the service. The IR only needs a stable list of -legal render targets and the small amount of metadata consumed by downstream -passes. +Legality has already happened in the service. Priority is consumed while +validating and ordering the wire response; it is not persisted. The IR only +needs a stable list of legal render targets and the small amount of metadata +consumed by downstream passes. Array position is meaningful: candidate zero is +the selected version. Do not add fields to the IR candidate payload simply because they exist in Python metadata. Add a field only when a C++ pass or IR-level test consumes it. @@ -202,7 +219,7 @@ must include every input that can change the rendered helper body: - op name - target architecture - tile operand dtype, shape, valid shape, memory space, layouts, fractal size, - and pad value + pad value, and compact mode - view operand dtype, shape, strides, memory space, and layout - vector operand dtype and shape - scalar operand dtype and static value when known @@ -220,6 +237,9 @@ faster to inspect. enough to accept ST-proven TileLangDSL forms. - Forward context attrs before porting a version that depends on them. - Use stable candidate ids for multi-candidate ops. +- Treat candidate ids as identity only; use priority for preference. +- Reject equal top-priority candidates instead of adding an incidental + tie-breaker. - Put all helper-code-affecting operand metadata in the specialization key. - Add a focused regression for each backend-selection bug. - Treat full ST status files as snapshots, not design documentation. diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 18cc73c8e8..f160a275e4 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -104,6 +104,7 @@ struct OperandTypeInfo { int32_t slayout = 0; int32_t fractal = 0; uint64_t pad = 0; + int32_t compact = 0; // --- View-only (MemRefType) — for JSON / constraint checking only --- SmallVector viewShape; @@ -126,7 +127,8 @@ struct OperandTypeInfo { tileValidShape == rhs.tileValidShape && tileMemorySpace == rhs.tileMemorySpace && blayout == rhs.blayout && slayout == rhs.slayout && - fractal == rhs.fractal && pad == rhs.pad; + fractal == rhs.fractal && pad == rhs.pad && + compact == rhs.compact; if (kind == OperandKind::Vector) return vectorShape == rhs.vectorShape; if (kind == OperandKind::Scalar) @@ -164,7 +166,7 @@ struct SpecKeyInfo : public llvm::DenseMapInfo { h = llvm::hash_combine(h, static_cast(op.kind), op.dtype); if (op.kind == OperandKind::Tile) { h = llvm::hash_combine(h, op.tileMemorySpace, op.blayout, - op.slayout, op.fractal, op.pad); + op.slayout, op.fractal, op.pad, op.compact); for (int64_t d : op.tileShape) h = llvm::hash_combine(h, d); for (int64_t d : op.tileValidShape) @@ -725,6 +727,8 @@ static std::optional buildOperandTypeInfo(Value value) { ? static_cast(config.getSFractalSize().getInt()) : 0; info.pad = static_cast(config.getPad().getValue()); + info.compact = + static_cast(config.getCompactMode().getValue()); } return info; } @@ -893,7 +897,9 @@ static std::string buildOperandSpecsJson(const SpecKey &key) { json += std::to_string(op.fractal); json += ",\"pad_value\":\"0x"; json += llvm::utohexstr(op.pad, /*LowerCase=*/false); - json += "\"}}"; + json += "\",\"compact_mode\":"; + json += std::to_string(op.compact); + json += "}}"; continue; } @@ -964,6 +970,7 @@ static std::string buildUniqueFunctionBaseName(const SpecKey &key) { uniqueName += "_sl" + std::to_string(op.slayout); uniqueName += "_fr" + std::to_string(op.fractal); uniqueName += "_pd" + llvm::utohexstr(op.pad, /*LowerCase=*/false); + uniqueName += "_cm" + std::to_string(op.compact); } else if (op.kind == OperandKind::View) { uniqueName += "_ms_" + op.viewMemorySpace; uniqueName += "_shape"; diff --git a/lib/PTO/Transforms/InsertTemplateAttributes.cpp b/lib/PTO/Transforms/InsertTemplateAttributes.cpp index 324d8c0a3f..eee8d3f879 100644 --- a/lib/PTO/Transforms/InsertTemplateAttributes.cpp +++ b/lib/PTO/Transforms/InsertTemplateAttributes.cpp @@ -19,6 +19,7 @@ #include "mlir/IR/BuiltinTypes.h" #include "mlir/Pass/Pass.h" +#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringExtras.h" @@ -47,6 +48,7 @@ constexpr llvm::StringLiteral kCandidatesAttr = "candidates"; struct CandidateMetadata { int64_t id; std::string name; + int64_t priority; int64_t loopDepth; bool postUpdate; bool tail; @@ -579,12 +581,14 @@ static void appendTileOperandSpecJson(std::string &json, pto::SLayout sLayout = pto::SLayout::NoneBox; int64_t fractalSize = 0; uint64_t padValue = 0; + int32_t compactMode = static_cast(pto::CompactMode::Null); if (auto config = tileType.getConfigAttr()) { bLayout = config.getBLayout().getValue(); sLayout = config.getSLayout().getValue(); if (config.getSFractalSize()) fractalSize = config.getSFractalSize().getInt(); padValue = static_cast(config.getPad().getValue()); + compactMode = static_cast(config.getCompactMode().getValue()); } json += "\",\"config\":{\"b_layout\":\""; @@ -595,7 +599,9 @@ static void appendTileOperandSpecJson(std::string &json, json += std::to_string(fractalSize); json += ",\"pad_value\":\"0x"; json += llvm::utohexstr(padValue, /*LowerCase=*/false); - json += "\"}}"; + json += "\",\"compact_mode\":"; + json += std::to_string(compactMode); + json += "}}"; } static void appendViewOperandSpecJson(std::string &json, Value operand, @@ -778,7 +784,7 @@ parseCandidateAttributes(Operation *operation, StringRef metadataJson) { } auto *root = parsed->getAsObject(); - auto *candidates = root ? root->getObject("candidates") : nullptr; + auto *candidates = root ? root->getArray("candidates") : nullptr; if (!candidates || candidates->empty()) { operation->emitError("InsertTemplateAttributes found no legal template " "candidates for ") @@ -788,8 +794,9 @@ parseCandidateAttributes(Operation *operation, StringRef metadataJson) { SmallVector parsedCandidates; parsedCandidates.reserve(candidates->size()); - for (const auto &entry : *candidates) { - auto *metadata = entry.second.getAsObject(); + llvm::DenseSet candidateIds; + for (const llvm::json::Value &entry : *candidates) { + auto *metadata = entry.getAsObject(); if (!metadata) { operation->emitError( "InsertTemplateAttributes candidate metadata must be an object"); @@ -798,13 +805,14 @@ parseCandidateAttributes(Operation *operation, StringRef metadataJson) { auto name = metadata->getString("name"); auto id = metadata->getInteger("id"); + auto priority = metadata->getInteger("priority"); auto loopDepth = metadata->getInteger("loop_depth"); auto postUpdate = metadata->getBoolean("is_post_update"); auto tail = metadata->getBoolean("has_tail"); - if (!name || !loopDepth || !postUpdate || !tail) { + if (!name || !priority || !loopDepth || !postUpdate || !tail) { operation->emitError( "InsertTemplateAttributes candidate metadata is missing name, " - "loop_depth, is_post_update, or has_tail"); + "priority, loop_depth, is_post_update, or has_tail"); return failure(); } if (!id && candidates->size() != 1) { @@ -814,9 +822,17 @@ parseCandidateAttributes(Operation *operation, StringRef metadataJson) { return failure(); } + int64_t candidateId = id.value_or(0); + if (!candidateIds.insert(candidateId).second) { + operation->emitError( + "InsertTemplateAttributes candidate ids must be unique"); + return failure(); + } + parsedCandidates.push_back(CandidateMetadata{ - id.value_or(0), + candidateId, name->str(), + *priority, *loopDepth, *postUpdate, *tail, @@ -826,16 +842,20 @@ parseCandidateAttributes(Operation *operation, StringRef metadataJson) { llvm::sort(parsedCandidates, [](const CandidateMetadata &left, const CandidateMetadata &right) { - if (left.id != right.id) - return left.id < right.id; + if (left.priority != right.priority) + return left.priority > right.priority; return left.name < right.name; }); - for (auto [index, candidate] : llvm::enumerate(parsedCandidates)) { - if (index != 0 && candidate.id == parsedCandidates[index - 1].id) { - operation->emitError( - "InsertTemplateAttributes candidate ids must be unique"); - return failure(); - } + if (parsedCandidates.size() > 1 && + parsedCandidates[0].priority == parsedCandidates[1].priority) { + operation->emitError( + "InsertTemplateAttributes found multiple legal templates tied at " + "the highest priority: ") + << parsedCandidates[0].name << " and " << parsedCandidates[1].name + << " at priority " << parsedCandidates[0].priority + << "; assign distinct priorities or make their constraints mutually " + "exclusive"; + return failure(); } Builder builder(operation->getContext()); diff --git a/lib/PTO/Transforms/TileFusion/PTOFusionPredicateElision.cpp b/lib/PTO/Transforms/TileFusion/PTOFusionPredicateElision.cpp index d5ec1f0df1..59ce4d1b60 100644 --- a/lib/PTO/Transforms/TileFusion/PTOFusionPredicateElision.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOFusionPredicateElision.cpp @@ -9,6 +9,7 @@ #include "PTO/IR/PTO.h" #include "PTO/Transforms/Passes.h" +#include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/Dominance.h" @@ -136,6 +137,12 @@ static std::optional getForIterArgInfo(Value value) { } static std::optional getPltScalarOutInfo(Value value) { + // Element-wise templates keep the remaining element count as index while + // plt consumes and returns an integer scalar. Look through the cast used to + // feed the scalar result back to scf.for. + if (auto indexCast = value.getDefiningOp()) + value = indexCast.getIn(); + auto result = dyn_cast(value); if (!result || result.getResultNumber() != 1) return std::nullopt; @@ -180,10 +187,18 @@ static bool areEquivalentLoopCarriedValues(Value lhs, Value rhs, if (lhsYieldInfo->bitWidth != rhsYieldInfo->bitWidth) return false; + Value lhsRecurrenceInput = lhsYieldInfo->scalar; + Value rhsRecurrenceInput = rhsYieldInfo->scalar; + if (auto indexCast = lhsRecurrenceInput.getDefiningOp()) + lhsRecurrenceInput = indexCast.getIn(); + if (auto indexCast = rhsRecurrenceInput.getDefiningOp()) + rhsRecurrenceInput = indexCast.getIn(); + // Stay conservative on unsupported cyclic proofs. The only accepted // recurrence cycle is the direct iter_arg -> plt.scalar_out self recursion - // for the same value pair; more complex cycles remain unsupported. - if (areSameValuePair(lhs, rhs, lhsYieldInfo->scalar, rhsYieldInfo->scalar)) + // for the same value pair, optionally bridged by the index casts required by + // the plt/scf type boundary; more complex cycles remain unsupported. + if (areSameValuePair(lhs, rhs, lhsRecurrenceInput, rhsRecurrenceInput)) return true; return areEquivalentValues(lhsYieldInfo->scalar, rhsYieldInfo->scalar, diff --git a/lib/TileOps/a5/_elementwise.py b/lib/TileOps/a5/_elementwise.py index e80ef495f5..b9d2c1d37e 100644 --- a/lib/TileOps/a5/_elementwise.py +++ b/lib/TileOps/a5/_elementwise.py @@ -5,12 +5,21 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""Shared PTODSL implementations for straightforward A5 elementwise TileOps.""" +"""Shared PTODSL implementations for straightforward A5 elementwise TileOps. + +The module-level traversal cores opt into PTODSL's source rewrite so they can +author runtime control flow with ordinary Python ``for range(...)`` syntax. +""" from ptodsl import pto +from ptodsl._ast_rewrite import rewrite_jit_function import ptodsl.tilelib as tilelib +FALLBACK_TRAVERSAL_PRIORITY = 0 +PREFERRED_TRAVERSAL_PRIORITY = 10 + + def _ub_or_vec_row_major(operand_memory_spaces, operand_b_layouts, operand_s_layouts, **_): return ( all(space in {"ub", "vec"} for space in operand_memory_spaces) @@ -26,10 +35,227 @@ def _common_constraints(*operand_names): ] -def register_unary(*, op, name, vector_op, dtypes, constraints=()): +def _traversal_loop_depth(traversal): + if traversal == "1d": + return 1 + if traversal == "2d": + return 2 + raise ValueError( + f"unsupported element-wise traversal {traversal!r}; expected '1d' or '2d'" + ) + + +def traversal_metadata( + traversal, + *, + priority=None, + candidate_id=None, + fallback_candidate_id=0, + candidate_count=1, +): + """Resolve stable candidate metadata for a 2D/1D traversal pair.""" + + loop_depth = _traversal_loop_depth(traversal) + if priority is None: + priority = ( + PREFERRED_TRAVERSAL_PRIORITY + if loop_depth == 1 + else FALLBACK_TRAVERSAL_PRIORITY + ) + if candidate_id is None: + candidate_id = fallback_candidate_id + if loop_depth == 1: + candidate_id += candidate_count + return loop_depth, priority, candidate_id + + +def _with_traversal_constraint(traversal, operand_names, constraints): + loop_depth = _traversal_loop_depth(traversal) + result = list(constraints) + if loop_depth == 1: + result.append(tilelib.require_elementwise_1d(*operand_names)) + return result + + +@rewrite_jit_function +def emit_elementwise_1d(anchor, emit_chunk): + """Traverse ``anchor`` as one contiguous logical element range. + + The caller is responsible for proving that every pointer used by + ``emit_chunk`` describes the same gap-free range. ``emit_chunk`` receives + the flattened element offset and the mask for that vector iteration. + """ + + dtype = anchor.dtype + valid_rows, valid_cols = anchor.valid_shape + lanes = pto.elements_per_vreg(dtype) + total_elements = valid_rows * valid_cols + remained = total_elements + for offset in range(0, total_elements, lanes): + mask, remained = pto.make_mask(dtype, remained) + emit_chunk(offset, mask) + + +@rewrite_jit_function +def emit_elementwise_2d(anchor, emit_chunk): + """Traverse ``anchor`` row by row with one vector loop per row. + + ``emit_chunk`` receives the row, column, and mask for the current vector + iteration. The per-row mask state is deliberately restarted at + ``valid_cols``. + """ + + dtype = anchor.dtype + valid_rows, valid_cols = anchor.valid_shape + lanes = pto.elements_per_vreg(dtype) + for row in range(0, valid_rows, 1): + remained = valid_cols + for col in range(0, valid_cols, lanes): + mask, remained = pto.make_mask(dtype, remained) + emit_chunk(row, col, mask) + + +def emit_unary_1d(src, dst, vector_op): + """Emit a flattened unary element-wise body.""" + + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + def emit_chunk(offset, mask): + value = pto.vlds(src_ptr, offset) + result = vector_op(value, mask) + pto.vsts(result, dst_ptr, offset, mask) + + emit_elementwise_1d(dst, emit_chunk) + + +def emit_unary_2d(src, dst, vector_op): + """Emit a row-wise unary element-wise body.""" + + def emit_chunk(row, col, mask): + value = pto.vlds(src[row, col:]) + result = vector_op(value, mask) + pto.vsts(result, dst[row, col:], mask) + + emit_elementwise_2d(dst, emit_chunk) + + +def emit_binary_1d(src0, src1, dst, vector_op): + """Emit a flattened Tile-Tile element-wise body.""" + + src0_ptr = src0.as_ptr() + src1_ptr = src1.as_ptr() + dst_ptr = dst.as_ptr() + + def emit_chunk(offset, mask): + lhs = pto.vlds(src0_ptr, offset) + rhs = pto.vlds(src1_ptr, offset) + result = vector_op(lhs, rhs, mask) + pto.vsts(result, dst_ptr, offset, mask) + + emit_elementwise_1d(dst, emit_chunk) + + +def emit_binary_2d(src0, src1, dst, vector_op): + """Emit a row-wise Tile-Tile element-wise body.""" + + def emit_chunk(row, col, mask): + lhs = pto.vlds(src0[row, col:]) + rhs = pto.vlds(src1[row, col:]) + result = vector_op(lhs, rhs, mask) + pto.vsts(result, dst[row, col:], mask) + + emit_elementwise_2d(dst, emit_chunk) + + +def _emit_scalar_compute(value, scalar, mask, vector_op, broadcast_scalar, + scalar_lhs): + if broadcast_scalar or scalar_lhs: + scalar_value = pto.vbr(scalar) + lhs, rhs = ( + (scalar_value, value) if scalar_lhs else (value, scalar_value) + ) + return vector_op(lhs, rhs, mask) + return vector_op(value, scalar, mask) + + +def emit_scalar_binary_1d(src, scalar, dst, vector_op, broadcast_scalar=False, + scalar_lhs=False): + """Emit a flattened Tile-Scalar or Scalar-Tile element-wise body.""" + + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + def emit_chunk(offset, mask): + value = pto.vlds(src_ptr, offset) + result = _emit_scalar_compute( + value, + scalar, + mask, + vector_op, + broadcast_scalar, + scalar_lhs, + ) + pto.vsts(result, dst_ptr, offset, mask) + + emit_elementwise_1d(dst, emit_chunk) + + +def emit_scalar_binary_2d(src, scalar, dst, vector_op, broadcast_scalar=False, + scalar_lhs=False): + """Emit a row-wise Tile-Scalar or Scalar-Tile element-wise body.""" + + def emit_chunk(row, col, mask): + value = pto.vlds(src[row, col:]) + result = _emit_scalar_compute( + value, + scalar, + mask, + vector_op, + broadcast_scalar, + scalar_lhs, + ) + pto.vsts(result, dst[row, col:], mask) + + emit_elementwise_2d(dst, emit_chunk) + + +def emit_scalar_fill_1d(scalar, dst): + """Emit a flattened scalar-fill body.""" + + dst_ptr = dst.as_ptr() + + def emit_chunk(offset, mask): + value = pto.vdup(scalar, mask) + pto.vsts(value, dst_ptr, offset, mask) + + emit_elementwise_1d(dst, emit_chunk) + + +def emit_scalar_fill_2d(scalar, dst): + """Emit a row-wise scalar-fill body.""" + + def emit_chunk(row, col, mask): + value = pto.vdup(scalar, mask) + pto.vsts(value, dst[row, col:], mask) + + emit_elementwise_2d(dst, emit_chunk) + + +def register_unary(*, op, name, vector_op, dtypes, constraints=(), + traversal="2d", priority=None, candidate_id=None): """Register a unary tile traversal using a public PTODSL vector operation.""" - candidate_constraints = _common_constraints("src", "dst") + list(constraints) + loop_depth, priority, candidate_id = traversal_metadata( + traversal, + priority=priority, + candidate_id=candidate_id, + ) + candidate_constraints = _with_traversal_constraint( + traversal, + ("src", "dst"), + _common_constraints("src", "dst") + list(constraints), + ) @tilelib.tile_template( op=op, @@ -40,31 +266,36 @@ def register_unary(*, op, name, vector_op, dtypes, constraints=()): op_engine="vector", op_class="elementwise", constraints=candidate_constraints, - id=0, - loop_depth=2, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, is_post_update=False, tags=("elementwise", "unary"), ) def template(src: pto.Tile, dst: pto.Tile): - dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - lanes = pto.elements_per_vreg(dtype) - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - value = pto.vlds(src[row, col:]) - result = vector_op(value, mask) - pto.vsts(result, dst[row, col:], mask) + if traversal == "1d": + emit_unary_1d(src, dst, vector_op) + else: + emit_unary_2d(src, dst, vector_op) return template -def register_binary(*, op, name, vector_op, dtypes, has_tmp=False): +def register_binary(*, op, name, vector_op, dtypes, has_tmp=False, + traversal="2d", priority=None, candidate_id=None): """Register a binary tile traversal, retaining an optional TileOp tmp operand.""" + loop_depth, priority, candidate_id = traversal_metadata( + traversal, + priority=priority, + candidate_id=candidate_id, + ) if has_tmp: + candidate_constraints = _with_traversal_constraint( + traversal, + ("src0", "src1", "tmp", "dst"), + _common_constraints("src0", "src1", "tmp", "dst"), + ) @tilelib.tile_template( op=op, @@ -74,29 +305,27 @@ def register_binary(*, op, name, vector_op, dtypes, has_tmp=False): iteration_axis="none", op_engine="vector", op_class="elementwise", - constraints=_common_constraints("src0", "src1", "tmp", "dst"), - id=0, - loop_depth=2, + constraints=candidate_constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, is_post_update=False, tags=("elementwise", "binary"), ) def template(src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, dst: pto.Tile): _ = tmp - dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - lanes = pto.elements_per_vreg(dtype) - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - result = vector_op(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) + if traversal == "1d": + emit_binary_1d(src0, src1, dst, vector_op) + else: + emit_binary_2d(src0, src1, dst, vector_op) return template + candidate_constraints = _with_traversal_constraint( + traversal, + ("src0", "src1", "dst"), + _common_constraints("src0", "src1", "dst"), + ) @tilelib.tile_template( op=op, target="a5", @@ -105,34 +334,36 @@ def template(src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, dst: pto.Tile): iteration_axis="none", op_engine="vector", op_class="elementwise", - constraints=_common_constraints("src0", "src1", "dst"), - id=0, - loop_depth=2, + constraints=candidate_constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, is_post_update=False, tags=("elementwise", "binary"), ) def template(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - lanes = pto.elements_per_vreg(dtype) - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - result = vector_op(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) + if traversal == "1d": + emit_binary_1d(src0, src1, dst, vector_op) + else: + emit_binary_2d(src0, src1, dst, vector_op) return template def register_scalar_binary(*, op, name, vector_op, dtypes, broadcast_scalar=False, has_tmp=False, tmp_matches_src_dst=True, - reverse_name=None): + reverse_name=None, traversal="2d", priority=None, + candidate_id=None, reverse_candidate_id=None): """Register a tile/scalar traversal using either a vector-scalar or broadcast op.""" + loop_depth, priority, candidate_id = traversal_metadata( + traversal, + priority=priority, + candidate_id=candidate_id, + candidate_count=2 if reverse_name else 1, + ) + if reverse_candidate_id is None: + reverse_candidate_id = candidate_id + 1 constraints = _common_constraints("src", "dst") if has_tmp: if tmp_matches_src_dst: @@ -141,6 +372,11 @@ def register_scalar_binary(*, op, name, vector_op, dtypes, broadcast_scalar=Fals constraints = _common_constraints("src", "dst") + [ _ub_or_vec_row_major, ] + constraints = _with_traversal_constraint( + traversal, + ("src", "tmp", "dst"), + constraints, + ) @tilelib.tile_template( op=op, @@ -151,14 +387,30 @@ def register_scalar_binary(*, op, name, vector_op, dtypes, broadcast_scalar=Fals op_engine="vector", op_class="elementwise", constraints=constraints, - id=0, - loop_depth=2, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, is_post_update=False, tags=("elementwise", "scalar"), ) def template(src: pto.Tile, scalar, tmp: pto.Tile, dst: pto.Tile): _ = tmp - _emit_scalar_binary_body(src, scalar, dst, vector_op, broadcast_scalar) + if traversal == "1d": + emit_scalar_binary_1d( + src, + scalar, + dst, + vector_op, + broadcast_scalar, + ) + else: + emit_scalar_binary_2d( + src, + scalar, + dst, + vector_op, + broadcast_scalar, + ) if reverse_name: @@ -171,26 +423,42 @@ def template(src: pto.Tile, scalar, tmp: pto.Tile, dst: pto.Tile): op_engine="vector", op_class="elementwise", constraints=constraints, - id=1, - loop_depth=2, + priority=priority, + id=reverse_candidate_id, + loop_depth=loop_depth, is_post_update=False, tags=("elementwise", "scalar"), ) def reverse_template(scalar, src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): _ = tmp - _emit_scalar_binary_body( - src, - scalar, - dst, - vector_op, - broadcast_scalar=True, - scalar_lhs=True, - ) + if traversal == "1d": + emit_scalar_binary_1d( + src, + scalar, + dst, + vector_op, + broadcast_scalar=True, + scalar_lhs=True, + ) + else: + emit_scalar_binary_2d( + src, + scalar, + dst, + vector_op, + broadcast_scalar=True, + scalar_lhs=True, + ) return template, reverse_template return template + constraints = _with_traversal_constraint( + traversal, + ("src", "dst"), + constraints, + ) @tilelib.tile_template( op=op, target="a5", @@ -200,13 +468,29 @@ def reverse_template(scalar, src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): op_engine="vector", op_class="elementwise", constraints=constraints, - id=0, - loop_depth=2, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, is_post_update=False, tags=("elementwise", "scalar"), ) def template(src: pto.Tile, scalar, dst: pto.Tile): - _emit_scalar_binary_body(src, scalar, dst, vector_op, broadcast_scalar) + if traversal == "1d": + emit_scalar_binary_1d( + src, + scalar, + dst, + vector_op, + broadcast_scalar, + ) + else: + emit_scalar_binary_2d( + src, + scalar, + dst, + vector_op, + broadcast_scalar, + ) if reverse_name: @@ -219,29 +503,56 @@ def template(src: pto.Tile, scalar, dst: pto.Tile): op_engine="vector", op_class="elementwise", constraints=constraints, - id=1, - loop_depth=2, + priority=priority, + id=reverse_candidate_id, + loop_depth=loop_depth, is_post_update=False, tags=("elementwise", "scalar"), ) def reverse_template(scalar, src: pto.Tile, dst: pto.Tile): - _emit_scalar_binary_body( - src, - scalar, - dst, - vector_op, - broadcast_scalar=True, - scalar_lhs=True, - ) + if traversal == "1d": + emit_scalar_binary_1d( + src, + scalar, + dst, + vector_op, + broadcast_scalar=True, + scalar_lhs=True, + ) + else: + emit_scalar_binary_2d( + src, + scalar, + dst, + vector_op, + broadcast_scalar=True, + scalar_lhs=True, + ) return template, reverse_template return template -def register_scalar_fill(*, op, name, dtypes): +def register_scalar_fill(*, op, name, dtypes, traversal="2d", priority=None, + candidate_id=None): """Register a scalar-to-tile fill traversal.""" + loop_depth, priority, candidate_id = traversal_metadata( + traversal, + priority=priority, + candidate_id=candidate_id, + ) + constraints = [ + tilelib.check_memory_space("ub"), + tilelib.check_layout("row_major"), + tilelib.check_s_layout("none_box"), + ] + constraints = _with_traversal_constraint( + traversal, + ("dst",), + constraints, + ) @tilelib.tile_template( op=op, target="a5", @@ -250,61 +561,38 @@ def register_scalar_fill(*, op, name, dtypes): iteration_axis="none", op_engine="vector", op_class="elementwise", - constraints=[ - tilelib.check_memory_space("ub"), - tilelib.check_layout("row_major"), - tilelib.check_s_layout("none_box"), - ], - id=0, - loop_depth=2, + constraints=constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, is_post_update=False, tags=("elementwise", "scalar", "fill"), ) def template(scalar, dst: pto.Tile): - dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - cols = dst.shape[1] - lanes = pto.elements_per_vreg(dtype) - dst_ptr = dst.as_ptr() - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - value = pto.vdup(scalar, mask) - addr = pto.addptr(dst_ptr, row * cols + col) - pto.vsts(value, addr, 0, mask) + if traversal == "1d": + emit_scalar_fill_1d(scalar, dst) + else: + emit_scalar_fill_2d(scalar, dst) return template -def _emit_scalar_binary_body(src, scalar, dst, vector_op, broadcast_scalar, - scalar_lhs=False): - dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - lanes = pto.elements_per_vreg(dtype) - - with pto.for_(0, valid_rows, step=1) as row: - col_loop = pto.for_(0, valid_cols, step=lanes).carry(remained=valid_cols) - with col_loop: - col = col_loop.iv - mask, remained = pto.make_mask(dtype, col_loop.remained) - value = pto.vlds(src[row, col:]) - if broadcast_scalar or scalar_lhs: - scalar_value = pto.vbr(scalar) - lhs, rhs = ( - (scalar_value, value) if scalar_lhs else (value, scalar_value) - ) - result = vector_op(lhs, rhs, mask) - else: - result = vector_op(value, scalar, mask) - pto.vsts(result, dst[row, col:], mask) - col_loop.update(remained=remained) - - __all__ = [ + "FALLBACK_TRAVERSAL_PRIORITY", + "PREFERRED_TRAVERSAL_PRIORITY", + "emit_binary_1d", + "emit_binary_2d", + "emit_elementwise_1d", + "emit_elementwise_2d", + "emit_scalar_binary_1d", + "emit_scalar_binary_2d", + "emit_scalar_fill_1d", + "emit_scalar_fill_2d", + "emit_unary_1d", + "emit_unary_2d", "register_binary", "register_scalar_binary", "register_scalar_fill", "register_unary", + "traversal_metadata", ] diff --git a/lib/TileOps/a5/_remainder.py b/lib/TileOps/a5/_remainder.py index 0647856e16..db00bc05c4 100644 --- a/lib/TileOps/a5/_remainder.py +++ b/lib/TileOps/a5/_remainder.py @@ -11,6 +11,13 @@ import ptodsl.tilelib as tilelib from ._common import ub_row_major_constraints +from ._elementwise import ( + emit_binary_1d, + emit_binary_2d, + emit_scalar_binary_1d, + emit_scalar_binary_2d, + traversal_metadata, +) FMOD_DTYPES = [("f32", "f32", "f32"), ("f16", "f16", "f16"), ("i16", "i16", "i16"), ("ui16", "ui16", "ui16")] @@ -43,10 +50,31 @@ def _scalar_remainder(lhs, scalar, mask, *, round_mode, dtype): return pto.vsub(lhs, product, mask) -def register_binary_remainder(*, op, name, dtypes, round_mode, has_tmp=False): +def register_binary_remainder(*, op, name, dtypes, round_mode, has_tmp=False, + traversal="2d", priority=None, + candidate_id=None): + if traversal not in {"1d", "2d"}: + raise ValueError( + f"unsupported remainder traversal {traversal!r}; " + "expected '1d' or '2d'" + ) + loop_depth, priority, candidate_id = traversal_metadata( + traversal, + priority=priority, + candidate_id=candidate_id, + ) constraints = ub_row_major_constraints("src0", "src1", "dst") if has_tmp: constraints = ub_row_major_constraints("src0", "src1", "tmp", "dst") + if traversal == "1d": + constraints.append( + tilelib.require_elementwise_1d( + "src0", + "src1", + "tmp", + "dst", + ) + ) @tilelib.tile_template( op=op, @@ -57,17 +85,23 @@ def register_binary_remainder(*, op, name, dtypes, round_mode, has_tmp=False): op_engine="vector", op_class="elementwise", constraints=constraints, - id=0, - loop_depth=2, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, is_post_update=False, tags=("elementwise", "remainder"), ) def template(src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, dst: pto.Tile): _ = tmp - _emit_binary(src0, src1, dst, round_mode) + _emit_binary(src0, src1, dst, round_mode, traversal) return template + if traversal == "1d": + constraints.append( + tilelib.require_elementwise_1d("src0", "src1", "dst") + ) + @tilelib.tile_template( op=op, target="a5", @@ -77,37 +111,56 @@ def template(src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, dst: pto.Tile): op_engine="vector", op_class="elementwise", constraints=constraints, - id=0, - loop_depth=2, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, is_post_update=False, tags=("elementwise", "remainder"), ) def template(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - _emit_binary(src0, src1, dst, round_mode) + _emit_binary(src0, src1, dst, round_mode, traversal) return template -def _emit_binary(src0, src1, dst, round_mode): +def _emit_binary(src0, src1, dst, round_mode, traversal): dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - lanes = pto.elements_per_vreg(dtype) - with pto.for_(0, valid_rows, step=1) as row: - col_loop = pto.for_(0, valid_cols, step=lanes).carry(remained=valid_cols) - with col_loop: - col = col_loop.iv - mask, remained = pto.make_mask(dtype, col_loop.remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - result = _remainder(lhs, rhs, mask, round_mode=round_mode, dtype=dtype) - pto.vsts(result, dst[row, col:], mask) - col_loop.update(remained=remained) - - -def register_scalar_remainder(*, op, name, dtypes, round_mode, has_tmp=False): + + def remainder(lhs, rhs, mask): + return _remainder( + lhs, + rhs, + mask, + round_mode=round_mode, + dtype=dtype, + ) + + if traversal == "1d": + emit_binary_1d(src0, src1, dst, remainder) + else: + emit_binary_2d(src0, src1, dst, remainder) + + +def register_scalar_remainder(*, op, name, dtypes, round_mode, has_tmp=False, + traversal="2d", priority=None, + candidate_id=None): + if traversal not in {"1d", "2d"}: + raise ValueError( + f"unsupported remainder traversal {traversal!r}; " + "expected '1d' or '2d'" + ) + loop_depth, priority, candidate_id = traversal_metadata( + traversal, + priority=priority, + candidate_id=candidate_id, + ) constraints = ub_row_major_constraints("src", "dst") if has_tmp: constraints = ub_row_major_constraints("src", "tmp", "dst") + if traversal == "1d": + constraints.append( + tilelib.require_elementwise_1d("src", "tmp", "dst") + ) @tilelib.tile_template( op=op, @@ -118,17 +171,21 @@ def register_scalar_remainder(*, op, name, dtypes, round_mode, has_tmp=False): op_engine="vector", op_class="elementwise", constraints=constraints, - id=0, - loop_depth=2, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, is_post_update=False, tags=("elementwise", "scalar", "remainder"), ) def template(src: pto.Tile, scalar, tmp: pto.Tile, dst: pto.Tile): _ = tmp - _emit_scalar(src, scalar, dst, round_mode) + _emit_scalar(src, scalar, dst, round_mode, traversal) return template + if traversal == "1d": + constraints.append(tilelib.require_elementwise_1d("src", "dst")) + @tilelib.tile_template( op=op, target="a5", @@ -138,30 +195,34 @@ def template(src: pto.Tile, scalar, tmp: pto.Tile, dst: pto.Tile): op_engine="vector", op_class="elementwise", constraints=constraints, - id=0, - loop_depth=2, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, is_post_update=False, tags=("elementwise", "scalar", "remainder"), ) def template(src: pto.Tile, scalar, dst: pto.Tile): - _emit_scalar(src, scalar, dst, round_mode) + _emit_scalar(src, scalar, dst, round_mode, traversal) return template -def _emit_scalar(src, scalar, dst, round_mode): +def _emit_scalar(src, scalar, dst, round_mode, traversal): dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - lanes = pto.elements_per_vreg(dtype) - with pto.for_(0, valid_rows, step=1) as row: - col_loop = pto.for_(0, valid_cols, step=lanes).carry(remained=valid_cols) - with col_loop: - col = col_loop.iv - mask, remained = pto.make_mask(dtype, col_loop.remained) - value = pto.vlds(src[row, col:]) - result = _scalar_remainder(value, scalar, mask, round_mode=round_mode, dtype=dtype) - pto.vsts(result, dst[row, col:], mask) - col_loop.update(remained=remained) + + def remainder(value, scalar_value, mask): + return _scalar_remainder( + value, + scalar_value, + mask, + round_mode=round_mode, + dtype=dtype, + ) + + if traversal == "1d": + emit_scalar_binary_1d(src, scalar, dst, remainder) + else: + emit_scalar_binary_2d(src, scalar, dst, remainder) __all__ = [ diff --git a/lib/TileOps/a5/tabs.py b/lib/TileOps/a5/tabs.py index f3d0be67eb..ef6264c232 100644 --- a/lib/TileOps/a5/tabs.py +++ b/lib/TileOps/a5/tabs.py @@ -12,9 +12,21 @@ from ._elementwise import register_unary +_DTYPES = [("f16", "f16"), ("f32", "f32")] + + template_tabs = register_unary( op="pto.tabs", name="template_tabs", vector_op=pto.vabs, - dtypes=[("f16", "f16"), ("f32", "f32")], + dtypes=_DTYPES, +) + + +template_tabs_1d = register_unary( + op="pto.tabs", + name="template_tabs_1d", + vector_op=pto.vabs, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tadd.py b/lib/TileOps/a5/tadd.py index 2218a5d628..3b12331a96 100644 --- a/lib/TileOps/a5/tadd.py +++ b/lib/TileOps/a5/tadd.py @@ -17,9 +17,21 @@ def _vadd(lhs, rhs, mask): return pto.vadd(lhs, rhs, mask) +_DTYPES = same_dtype_signatures(3) + + template_tadd = register_binary( op="pto.tadd", name="template_tadd", vector_op=_vadd, - dtypes=same_dtype_signatures(3), + dtypes=_DTYPES, +) + + +template_tadd_1d = register_binary( + op="pto.tadd", + name="template_tadd_1d", + vector_op=_vadd, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tadds.py b/lib/TileOps/a5/tadds.py index eb77c11b27..953460b75f 100644 --- a/lib/TileOps/a5/tadds.py +++ b/lib/TileOps/a5/tadds.py @@ -13,9 +13,20 @@ from ._elementwise import register_scalar_binary +_DTYPES = same_dtype_signatures(3) + + template_tadds = register_scalar_binary( op="pto.tadds", name="template_tadds", vector_op=pto.vadds, - dtypes=same_dtype_signatures(3), + dtypes=_DTYPES, +) + +template_tadds_1d = register_scalar_binary( + op="pto.tadds", + name="template_tadds_1d", + vector_op=pto.vadds, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tand.py b/lib/TileOps/a5/tand.py index 992b19523b..c0a440fed9 100644 --- a/lib/TileOps/a5/tand.py +++ b/lib/TileOps/a5/tand.py @@ -13,9 +13,21 @@ from ._elementwise import register_binary +_DTYPES = [(dtype, dtype, dtype) for dtype in INT_DTYPES] + + template_tand = register_binary( op="pto.tand", name="template_tand", vector_op=pto.vand, - dtypes=[(dtype, dtype, dtype) for dtype in INT_DTYPES], + dtypes=_DTYPES, +) + + +template_tand_1d = register_binary( + op="pto.tand", + name="template_tand_1d", + vector_op=pto.vand, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tands.py b/lib/TileOps/a5/tands.py index 114ac78aab..e41f6171cc 100644 --- a/lib/TileOps/a5/tands.py +++ b/lib/TileOps/a5/tands.py @@ -13,10 +13,22 @@ from ._elementwise import register_scalar_binary +_DTYPES = [(dtype, dtype, dtype) for dtype in INT_DTYPES] + + template_tands = register_scalar_binary( op="pto.tands", name="template_tands", vector_op=pto.vand, broadcast_scalar=True, - dtypes=[(dtype, dtype, dtype) for dtype in INT_DTYPES], + dtypes=_DTYPES, +) + +template_tands_1d = register_scalar_binary( + op="pto.tands", + name="template_tands_1d", + vector_op=pto.vand, + broadcast_scalar=True, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tcmp.py b/lib/TileOps/a5/tcmp.py index f049014b41..3eb355687c 100644 --- a/lib/TileOps/a5/tcmp.py +++ b/lib/TileOps/a5/tcmp.py @@ -5,12 +5,31 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""PTODSL TileLib template for pto.tcmp.""" +"""PTODSL TileLib templates for pto.tcmp.""" from ptodsl import pto +from ptodsl._ast_rewrite import rewrite_jit_function import ptodsl.tilelib as tilelib -def _ub_or_vec_row_major(operand_memory_spaces, operand_b_layouts, operand_s_layouts, **_): +from ._elementwise import traversal_metadata + + +_DTYPES = [ + ("f32", "f32", "i8"), + ("i32", "i32", "i8"), + ("f16", "f16", "i8"), + ("i16", "i16", "i8"), + ("i8", "i8", "i8"), + ("ui8", "ui8", "i8"), +] + + +def _ub_or_vec_row_major( + operand_memory_spaces, + operand_b_layouts, + operand_s_layouts, + **_, +): return ( all(space in {"ub", "vec"} for space in operand_memory_spaces) and all(layout == "row_major" for layout in operand_b_layouts) @@ -18,31 +37,83 @@ def _ub_or_vec_row_major(operand_memory_spaces, operand_b_layouts, operand_s_lay ) -@tilelib.tile_template( - op="pto.tcmp", - target="a5", - name="template_tcmp", - dtypes=[ - ("f32", "f32", "i8"), - ("i32", "i32", "i8"), - ("f16", "f16", "i8"), - ("i16", "i16", "i8"), - ("i8", "i8", "i8"), - ("ui8", "ui8", "i8"), - ], - iteration_axis="none", - op_engine="vector", - op_class="other", - constraints=[ - _ub_or_vec_row_major, - tilelib.require_same_valid_shape("src0", "src1", "dst"), - ], - id=0, - loop_depth=2, - is_post_update=False, - tags=("compare", "predicate-store"), -) -def template_tcmp(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): +@rewrite_jit_function +def _emit_tcmp_1d(src0, src1, dst): + dtype = src0.dtype + valid_rows, valid_cols = src0.valid_shape + lanes = pto.elements_per_vreg(dtype) + total_elements = valid_rows * valid_cols + cmp_mode = pto.get_op_attr("cmp_mode", "eq") + src0_ptr = src0.as_ptr() + src1_ptr = src1.as_ptr() + dst_ptr = dst.as_ptr() + + if str(dtype) in {"f32", "i32"}: + remained = total_elements + for offset in range(0, total_elements, lanes * 2): + first_mask, remained = pto.make_mask(dtype, remained) + first_lhs = pto.vlds(src0_ptr, offset) + first_rhs = pto.vlds(src1_ptr, offset) + first_cmp = pto.vcmp( + first_lhs, + first_rhs, + first_mask, + cmp_mode, + ) + first_cmp_b8 = pto.pbitcast(first_cmp, pto.mask_b8) + + second_mask, remained = pto.make_mask(dtype, remained) + second_lhs = pto.vlds(src0_ptr, offset + lanes) + second_rhs = pto.vlds(src1_ptr, offset + lanes) + second_cmp = pto.vcmp( + second_lhs, + second_rhs, + second_mask, + cmp_mode, + ) + second_cmp_b8 = pto.pbitcast(second_cmp, pto.mask_b8) + + packed_low, _ = pto.pdintlv_b8( + first_cmp_b8, + second_cmp_b8, + ) + pto.psts( + packed_low, + dst_ptr, + offset // 8, + dist=pto.PredicateDist.PK, + ) + elif str(dtype) in {"f16", "i16"}: + remained = total_elements + for offset in range(0, total_elements, lanes): + mask, remained = pto.make_mask(dtype, remained) + lhs = pto.vlds(src0_ptr, offset) + rhs = pto.vlds(src1_ptr, offset) + cmp = pto.vcmp(lhs, rhs, mask, cmp_mode) + cmp_b8 = pto.pbitcast(cmp, pto.mask_b8) + pto.psts( + cmp_b8, + dst_ptr, + offset // 8, + dist=pto.PredicateDist.PK, + ) + else: + remained = total_elements + for offset in range(0, total_elements, lanes): + mask, remained = pto.make_mask(dtype, remained) + lhs = pto.vlds(src0_ptr, offset) + rhs = pto.vlds(src1_ptr, offset) + cmp = pto.vcmp(lhs, rhs, mask, cmp_mode) + pto.psts( + cmp, + dst_ptr, + offset // 8, + dist=pto.PredicateDist.NORM, + ) + + +@rewrite_jit_function +def _emit_tcmp_2d(src0, src1, dst): dtype = src0.dtype valid_rows, valid_cols = src0.valid_shape lanes = pto.elements_per_vreg(dtype) @@ -53,7 +124,6 @@ def template_tcmp(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): if str(dtype) in {"f32", "i32"}: repeat_times = (valid_cols + lanes - 1) // lanes + 1 iterations = repeat_times // 2 - for row in range(0, valid_rows, 1): remained = valid_cols for col in range(0, iterations, 1): @@ -63,21 +133,38 @@ def template_tcmp(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): first_mask, remained = pto.make_mask(dtype, remained) first_lhs = pto.vlds(src0[row, first_offset:]) first_rhs = pto.vlds(src1[row, first_offset:]) - first_cmp = pto.vcmp(first_lhs, first_rhs, first_mask, cmp_mode) + first_cmp = pto.vcmp( + first_lhs, + first_rhs, + first_mask, + cmp_mode, + ) first_cmp_b8 = pto.pbitcast(first_cmp, pto.mask_b8) second_mask, remained = pto.make_mask(dtype, remained) second_lhs = pto.vlds(src0[row, second_offset:]) second_rhs = pto.vlds(src1[row, second_offset:]) - second_cmp = pto.vcmp(second_lhs, second_rhs, second_mask, cmp_mode) + second_cmp = pto.vcmp( + second_lhs, + second_rhs, + second_mask, + cmp_mode, + ) second_cmp_b8 = pto.pbitcast(second_cmp, pto.mask_b8) - packed_low, _ = pto.pdintlv_b8(first_cmp_b8, second_cmp_b8) + packed_low, _ = pto.pdintlv_b8( + first_cmp_b8, + second_cmp_b8, + ) store_offset = row * dst_stride + col * 16 - pto.psts(packed_low, dst_ptr, store_offset, dist=pto.PredicateDist.PK) + pto.psts( + packed_low, + dst_ptr, + store_offset, + dist=pto.PredicateDist.PK, + ) elif str(dtype) in {"f16", "i16"}: iterations = (valid_cols + lanes - 1) // lanes - for row in range(0, valid_rows, 1): remained = valid_cols for col in range(0, iterations, 1): @@ -87,10 +174,14 @@ def template_tcmp(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): cmp = pto.vcmp(lhs, rhs, mask, cmp_mode) cmp_b8 = pto.pbitcast(cmp, pto.mask_b8) store_offset = row * dst_stride + col * 16 - pto.psts(cmp_b8, dst_ptr, store_offset, dist=pto.PredicateDist.PK) + pto.psts( + cmp_b8, + dst_ptr, + store_offset, + dist=pto.PredicateDist.PK, + ) else: iterations = (valid_cols + lanes - 1) // lanes - for row in range(0, valid_rows, 1): remained = valid_cols for col in range(0, iterations, 1): @@ -99,4 +190,59 @@ def template_tcmp(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): rhs = pto.vlds(src1[row, col * lanes:]) cmp = pto.vcmp(lhs, rhs, mask, cmp_mode) store_offset = row * dst_stride + col * 32 - pto.psts(cmp, dst_ptr, store_offset, dist=pto.PredicateDist.NORM) + pto.psts( + cmp, + dst_ptr, + store_offset, + dist=pto.PredicateDist.NORM, + ) + + +def _register_tcmp(*, name, traversal): + constraints = [ + _ub_or_vec_row_major, + tilelib.require_same_valid_shape("src0", "src1", "dst"), + ] + loop_depth, priority, candidate_id = traversal_metadata(traversal) + if traversal == "1d": + constraints.append( + tilelib.require_predicate_compare_1d( + "src0", + "src1", + predicate_operand="dst", + ) + ) + + @tilelib.tile_template( + op="pto.tcmp", + target="a5", + name=name, + dtypes=_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="other", + constraints=constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, + is_post_update=False, + tags=("compare", "predicate-store"), + ) + def template(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): + if traversal == "1d": + _emit_tcmp_1d(src0, src1, dst) + else: + _emit_tcmp_2d(src0, src1, dst) + + return template + + +template_tcmp = _register_tcmp( + name="template_tcmp", + traversal="2d", +) + +template_tcmp_1d = _register_tcmp( + name="template_tcmp_1d", + traversal="1d", +) diff --git a/lib/TileOps/a5/tcmps.py b/lib/TileOps/a5/tcmps.py index 24253f6d31..71af0965ad 100644 --- a/lib/TileOps/a5/tcmps.py +++ b/lib/TileOps/a5/tcmps.py @@ -5,13 +5,31 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""PTODSL TileLib template for pto.tcmps.""" +"""PTODSL TileLib templates for pto.tcmps.""" from ptodsl import pto +from ptodsl._ast_rewrite import rewrite_jit_function import ptodsl.tilelib as tilelib +from ._elementwise import traversal_metadata -def _ub_or_vec_row_major(operand_memory_spaces, operand_b_layouts, operand_s_layouts, **_): + +_DTYPES = [ + ("f32", "f32", "ui8"), + ("i32", "i32", "ui8"), + ("f16", "f16", "ui8"), + ("i16", "i16", "ui8"), + ("i8", "i8", "ui8"), + ("ui8", "ui8", "ui8"), +] + + +def _ub_or_vec_row_major( + operand_memory_spaces, + operand_b_layouts, + operand_s_layouts, + **_, +): return ( all(space in {"ub", "vec"} for space in operand_memory_spaces) and all(layout == "row_major" for layout in operand_b_layouts) @@ -19,83 +37,203 @@ def _ub_or_vec_row_major(operand_memory_spaces, operand_b_layouts, operand_s_lay ) -@tilelib.tile_template( - op="pto.tcmps", - target="a5", - name="template_tcmps", - dtypes=[ - ("f32", "f32", "ui8"), - ("i32", "i32", "ui8"), - ("f16", "f16", "ui8"), - ("i16", "i16", "ui8"), - ("i8", "i8", "ui8"), - ("ui8", "ui8", "ui8"), - ], - iteration_axis="none", - op_engine="vector", - op_class="other", - constraints=[_ub_or_vec_row_major], - id=0, - loop_depth=2, - is_post_update=False, - tags=("compare", "scalar", "predicate-store"), -) -def template_tcmps(src: pto.Tile, scalar, dst: pto.Tile): +@rewrite_jit_function +def _emit_tcmps_1d(src, scalar, dst): dtype = src.dtype valid_rows, valid_cols = src.valid_shape lanes = pto.elements_per_vreg(dtype) + total_elements = valid_rows * valid_cols cmp_mode = pto.get_op_attr("cmp_mode", "eq") - dst_ptr = dst.as_ptr() src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() if str(dtype) in {"f32", "i32"}: - total_elm = valid_rows * valid_cols - repeat_times = (total_elm + lanes - 1) // lanes + 1 - iterations = repeat_times // 2 + remained = total_elements + for offset in range(0, total_elements, lanes * 2): + first_mask, remained = pto.make_mask(dtype, remained) + first = pto.vlds(src_ptr, offset) + first_cmp = pto.vcmps(first, scalar, first_mask, cmp_mode) + first_cmp_b8 = pto.pbitcast(first_cmp, pto.mask_b8) - for iteration in range(0, iterations, 1): - elem_offset0 = iteration * 2 * lanes - elem_offset1 = (iteration * 2 + 1) * lanes + second_mask, remained = pto.make_mask(dtype, remained) + second = pto.vlds(src_ptr, offset + lanes) + second_cmp = pto.vcmps( + second, + scalar, + second_mask, + cmp_mode, + ) + second_cmp_b8 = pto.pbitcast(second_cmp, pto.mask_b8) - remaining0 = total_elm - elem_offset0 - remaining1 = total_elm - elem_offset1 + packed_low, _ = pto.pdintlv_b8( + first_cmp_b8, + second_cmp_b8, + ) + pto.psts( + packed_low, + dst_ptr, + offset // 8, + dist=pto.PredicateDist.PK, + ) + elif str(dtype) in {"f16", "i16"}: + remained = total_elements + for offset in range(0, total_elements, lanes): + mask, remained = pto.make_mask(dtype, remained) + value = pto.vlds(src_ptr, offset) + cmp = pto.vcmps(value, scalar, mask, cmp_mode) + pto.psts( + cmp, + dst_ptr, + offset // 8, + dist=pto.PredicateDist.PK, + ) + else: + remained = total_elements + for offset in range(0, total_elements, lanes): + mask, remained = pto.make_mask(dtype, remained) + value = pto.vlds(src_ptr, offset) + cmp = pto.vcmps(value, scalar, mask, cmp_mode) + pto.psts( + cmp, + dst_ptr, + offset // 8, + dist=pto.PredicateDist.NORM, + ) - mask0, _ = pto.make_mask(dtype, remaining0) - mask1, _ = pto.make_mask(dtype, remaining1) - vec0 = pto.vlds(src_ptr, elem_offset0) - vec1 = pto.vlds(src_ptr, elem_offset1) +@rewrite_jit_function +def _emit_tcmps_2d(src, scalar, dst): + dtype = src.dtype + valid_rows, valid_cols = src.valid_shape + lanes = pto.elements_per_vreg(dtype) + cmp_mode = pto.get_op_attr("cmp_mode", "eq") + dst_ptr = dst.as_ptr() - cmp0 = pto.vcmps(vec0, scalar, mask0, cmp_mode) - cmp1 = pto.vcmps(vec1, scalar, mask1, cmp_mode) + # The destination is a linear packed-predicate byte stream. Its physical + # column count provides capacity; it is not the stride between logical + # predicate rows. Each row starts after the packed bytes emitted for the + # preceding row, preserving the existing TCMPS layout contract. + if str(dtype) in {"f32", "i32"}: + repeat_times = (valid_cols + lanes - 1) // lanes + 1 + iterations = repeat_times // 2 + packed_row_bytes = iterations * 16 + for row in range(0, valid_rows, 1): + remained = valid_cols + for col in range(0, iterations, 1): + first_offset = col * lanes * 2 + second_offset = (col * 2 + 1) * lanes - cmp0_b8 = pto.pbitcast(cmp0, pto.mask_b8) - cmp1_b8 = pto.pbitcast(cmp1, pto.mask_b8) - packed_low, _ = pto.pdintlv_b8(cmp0_b8, cmp1_b8) + first_mask, remained = pto.make_mask(dtype, remained) + first = pto.vlds(src[row, first_offset:]) + first_cmp = pto.vcmps( + first, + scalar, + first_mask, + cmp_mode, + ) + first_cmp_b8 = pto.pbitcast(first_cmp, pto.mask_b8) - store_offset = iteration * 16 - pto.psts(packed_low, dst_ptr, store_offset, dist=pto.PredicateDist.PK) + second_mask, remained = pto.make_mask(dtype, remained) + second = pto.vlds(src[row, second_offset:]) + second_cmp = pto.vcmps( + second, + scalar, + second_mask, + cmp_mode, + ) + second_cmp_b8 = pto.pbitcast(second_cmp, pto.mask_b8) + + packed_low, _ = pto.pdintlv_b8( + first_cmp_b8, + second_cmp_b8, + ) + store_offset = row * packed_row_bytes + col * 16 + pto.psts( + packed_low, + dst_ptr, + store_offset, + dist=pto.PredicateDist.PK, + ) elif str(dtype) in {"f16", "i16"}: bytes_per_iter = 16 iters_per_row = (valid_cols + lanes - 1) // lanes - for row in range(0, valid_rows, 1): remained = valid_cols for col in range(0, valid_cols, lanes): mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - cmp = pto.vcmps(vec, scalar, mask, cmp_mode) - store_offset = (row * iters_per_row + col // lanes) * bytes_per_iter - pto.psts(cmp, dst_ptr, store_offset, dist=pto.PredicateDist.PK) + value = pto.vlds(src[row, col:]) + cmp = pto.vcmps(value, scalar, mask, cmp_mode) + store_offset = ( + row * iters_per_row + col // lanes + ) * bytes_per_iter + pto.psts( + cmp, + dst_ptr, + store_offset, + dist=pto.PredicateDist.PK, + ) else: bytes_per_iter = 32 iters_per_row = (valid_cols + lanes - 1) // lanes - for row in range(0, valid_rows, 1): remained = valid_cols for col in range(0, valid_cols, lanes): mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - cmp = pto.vcmps(vec, scalar, mask, cmp_mode) - store_offset = (row * iters_per_row + col // lanes) * bytes_per_iter - pto.psts(cmp, dst_ptr, store_offset, dist=pto.PredicateDist.NORM) + value = pto.vlds(src[row, col:]) + cmp = pto.vcmps(value, scalar, mask, cmp_mode) + store_offset = ( + row * iters_per_row + col // lanes + ) * bytes_per_iter + pto.psts( + cmp, + dst_ptr, + store_offset, + dist=pto.PredicateDist.NORM, + ) + + +def _register_tcmps(*, name, traversal): + constraints = [_ub_or_vec_row_major] + loop_depth, priority, candidate_id = traversal_metadata(traversal) + if traversal == "1d": + constraints.append( + tilelib.require_predicate_compare_1d( + "src", + predicate_operand="dst", + flattened_destination=True, + ) + ) + + @tilelib.tile_template( + op="pto.tcmps", + target="a5", + name=name, + dtypes=_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="other", + constraints=constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, + is_post_update=False, + tags=("compare", "scalar", "predicate-store"), + ) + def template(src: pto.Tile, scalar, dst: pto.Tile): + if traversal == "1d": + _emit_tcmps_1d(src, scalar, dst) + else: + _emit_tcmps_2d(src, scalar, dst) + + return template + + +template_tcmps = _register_tcmps( + name="template_tcmps", + traversal="2d", +) + +template_tcmps_1d = _register_tcmps( + name="template_tcmps_1d", + traversal="1d", +) diff --git a/lib/TileOps/a5/tcvt.py b/lib/TileOps/a5/tcvt.py index e54a360d79..24812cc971 100644 --- a/lib/TileOps/a5/tcvt.py +++ b/lib/TileOps/a5/tcvt.py @@ -5,13 +5,34 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""PTODSL TileLib templates for row-wise ``pto.tcvt`` paths.""" +"""PTODSL TileLib templates for A5 ``pto.tcvt`` 1D/2D paths. + +Conversion-specific instruction sequences stay in this module. Flattened +forms use typed source and destination pointers so dtype width changes retain +their existing logical-element addressing and distribution modes. +""" + +from dataclasses import replace from ptodsl import pto +from ptodsl._ast_rewrite import rewrite_jit_function from ptodsl._surface_values import unwrap_surface_value, wrap_surface_value import ptodsl.tilelib as tilelib from ptoas.mlir.dialects import pto as _pto +from ._elementwise import ( + FALLBACK_TRAVERSAL_PRIORITY, + PREFERRED_TRAVERSAL_PRIORITY, +) + + +_PENDING_TCVT_1D = [] + + +def _defer_tcvt_1d(candidate): + _PENDING_TCVT_1D.append(candidate) + return candidate + def _rowwise(src_shape, src_valid_shape, dst_shape, dst_valid_shape, src_config, dst_config, **_): return ( @@ -75,6 +96,52 @@ def _vselr_low_precision(src, idx): return wrap_surface_value(_pto.VselrOp(raw_src.type, raw_src, unwrap_surface_value(idx)).result) +def _tcvt_conversion_mask(src, mask, mode): + if mode == "src_full": + return pto.make_mask(src.dtype, pto.PAT.ALL) + return mask + + +def _tcvt_load_2d(src, row, col, dist): + if dist: + return pto.vlds(src[row, col:], dist=dist) + return pto.vlds(src[row, col:]) + + +def _tcvt_convert(vec, dtype, mask, *, rnd, sat, part): + kwargs = {} + if rnd: + kwargs["rnd"] = _round_mode() + sat_mode = _sat_mode(sat) + if sat_mode is not None: + kwargs["sat"] = sat_mode + part_mode = _part_mode(part) + if part_mode is not None: + kwargs["part"] = part_mode + return pto.vcvt(vec, dtype, mask, **kwargs) + + +def _tcvt_store_2d(converted, dst, row, col, mask, dist): + if dist: + pto.vsts(converted, dst[row, col:], mask, dist=dist) + else: + pto.vsts(converted, dst[row, col:], mask) + + +@rewrite_jit_function +def _emit_tcvt_1d(dst, step_dtype, mask_dtype, remaining_scale, emit_chunk): + """Run one conversion vector loop over the flattened destination range.""" + + valid_rows, valid_cols = dst.valid_shape + total_elements = valid_rows * valid_cols + lanes = pto.elements_per_vreg(step_dtype) + remained = total_elements * remaining_scale + for offset in range(0, total_elements, lanes): + mask, remained = pto.make_mask(mask_dtype, remained) + emit_chunk(offset, mask) + + +@rewrite_jit_function def _render_tcvt( src, dst, @@ -91,31 +158,76 @@ def _render_tcvt( dtype = dst.dtype loop_dtype = src.dtype if mask_dtype == "src" else dtype lanes = pto.elements_per_vreg(loop_dtype) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols - col_loop = pto.for_(0, valid_cols, step=lanes).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes): mask, remained = pto.make_mask(loop_dtype, remained) - convert_mask_value = mask - if convert_mask == "src_full": - convert_mask_value = pto.make_mask(src.dtype, pto.PAT.ALL) - vec = pto.vlds(src[row, col:], dist=load_dist) if load_dist else pto.vlds(src[row, col:]) - kwargs = {} - if rnd: - kwargs["rnd"] = _round_mode() - sat_mode = _sat_mode(sat) - if sat_mode is not None: - kwargs["sat"] = sat_mode - part_mode = _part_mode(part) - if part_mode is not None: - kwargs["part"] = part_mode - converted = pto.vcvt(vec, dtype, convert_mask_value, **kwargs) - if store_dist: - pto.vsts(converted, dst[row, col:], mask, dist=store_dist) - else: - pto.vsts(converted, dst[row, col:], mask) - col_loop.update(remained=remained) + convert_mask_value = _tcvt_conversion_mask( + src, + mask, + convert_mask, + ) + vec = _tcvt_load_2d(src, row, col, load_dist) + converted = _tcvt_convert( + vec, + dtype, + convert_mask_value, + rnd=rnd, + sat=sat, + part=part, + ) + _tcvt_store_2d( + converted, + dst, + row, + col, + mask, + store_dist, + ) + + +def _render_tcvt_1d( + src, + dst, + *, + rnd=False, + sat=None, + part=None, + load_dist=None, + store_dist=None, + mask_dtype="dst", + convert_mask="store", +): + dtype = dst.dtype + loop_dtype = src.dtype if mask_dtype == "src" else dtype + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + def emit_chunk(offset, mask): + convert_mask_value = mask + if convert_mask == "src_full": + convert_mask_value = pto.make_mask(src.dtype, pto.PAT.ALL) + vec = ( + pto.vlds(src_ptr, offset, dist=load_dist) + if load_dist + else pto.vlds(src_ptr, offset) + ) + kwargs = {} + if rnd: + kwargs["rnd"] = _round_mode() + sat_mode = _sat_mode(sat) + if sat_mode is not None: + kwargs["sat"] = sat_mode + part_mode = _part_mode(part) + if part_mode is not None: + kwargs["part"] = part_mode + converted = pto.vcvt(vec, dtype, convert_mask_value, **kwargs) + if store_dist: + pto.vsts(converted, dst_ptr, offset, mask, dist=store_dist) + else: + pto.vsts(converted, dst_ptr, offset, mask) + + _emit_tcvt_1d(dst, loop_dtype, loop_dtype, 1, emit_chunk) def _register_tcvt( @@ -131,34 +243,111 @@ def _register_tcvt( mask_dtype="dst", convert_mask="store", ): + def register_form( + *, traversal, candidate_name, candidate_id, priority, register=True + ): + constraints = [_rowwise] + if traversal == "1d": + constraints.append(tilelib.require_conversion_1d()) + + @tilelib.tile_template( + op="pto.tcvt", + target="a5", + name=candidate_name, + dtypes=[dtypes], + iteration_axis="none", + op_engine="vector", + op_class="other", + constraints=constraints, + priority=priority, + id=candidate_id, + loop_depth=1 if traversal == "1d" else 2, + is_post_update=False, + tags=("convert", traversal), + register=register, + ) + def template(src: pto.Tile, dst: pto.Tile): + renderer = _render_tcvt_1d if traversal == "1d" else _render_tcvt + renderer( + src, + dst, + rnd=rnd, + sat=sat, + part=part, + load_dist=load_dist, + store_dist=store_dist, + mask_dtype=mask_dtype, + convert_mask=convert_mask, + ) + + return template + + fallback = register_form( + traversal="2d", + candidate_name=name, + candidate_id=idx, + priority=FALLBACK_TRAVERSAL_PRIORITY, + ) + _defer_tcvt_1d( + register_form( + traversal="1d", + candidate_name=f"{name}_1d", + candidate_id=None, + priority=PREFERRED_TRAVERSAL_PRIORITY, + register=False, + ) + ) + return fallback + + +def _register_tcvt_1d( + *, + name, + dtypes, + renderer, + source_elements_per_destination=1, + tags=(), +): + """Register a preferred flattened form for a bespoke conversion body.""" + + dtype_signatures = ( + [dtypes] + if dtypes and isinstance(dtypes[0], str) + else list(dtypes) + ) + shape_constraint = ( + _rowwise_bf16_to_fp4 + if source_elements_per_destination == 2 + else _rowwise + ) + @tilelib.tile_template( op="pto.tcvt", target="a5", - name=name, - dtypes=[dtypes], + name=f"{name}_1d", + dtypes=dtype_signatures, iteration_axis="none", op_engine="vector", op_class="other", - constraints=[_rowwise], - id=idx, - loop_depth=2, + constraints=[ + shape_constraint, + tilelib.require_conversion_1d( + source_elements_per_destination=( + source_elements_per_destination + ), + ), + ], + priority=PREFERRED_TRAVERSAL_PRIORITY, + id=None, + loop_depth=1, is_post_update=False, - tags=("convert", "rowwise"), + tags=("convert", "1d", *tags), + register=False, ) def template(src: pto.Tile, dst: pto.Tile): - _render_tcvt( - src, - dst, - rnd=rnd, - sat=sat, - part=part, - load_dist=load_dist, - store_dist=store_dist, - mask_dtype=mask_dtype, - convert_mask=convert_mask, - ) + renderer(src, dst) - return template + return _defer_tcvt_1d(template) template_tcvt_f32_to_i32 = _register_tcvt( @@ -202,11 +391,9 @@ def template_tcvt_f16_to_i16(src: pto.Tile, dst: pto.Tile): lanes_f32 = pto.elements_per_vreg(pto.f32) full_mask_b16 = pto.make_mask(src.dtype, pto.PAT.ALL) full_mask_b32 = pto.make_mask(pto.i32, pto.PAT.ALL) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols - col_loop = pto.for_(0, valid_cols, step=lanes_f32).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_f32): store_mask, remained = pto.make_mask(pto.i32, remained) vec_f16 = pto.vlds(src[row, col:], dist="UNPK_B16") vec_i32 = pto.vcvt( @@ -224,7 +411,46 @@ def template_tcvt_f16_to_i16(src: pto.Tile, dst: pto.Tile): part=pto.VcvtPartMode.EVEN, ) pto.vsts(vec_i16, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B32) - col_loop.update(remained=remained) + + +def _render_tcvt_f16_to_i16_1d(src: pto.Tile, dst: pto.Tile): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + full_mask_b16 = pto.make_mask(src.dtype, pto.PAT.ALL) + full_mask_b32 = pto.make_mask(pto.i32, pto.PAT.ALL) + + def emit_chunk(offset, store_mask): + vec_f16 = pto.vlds(src_ptr, offset, dist="UNPK_B16") + vec_i32 = pto.vcvt( + vec_f16, + pto.i32, + full_mask_b16, + rnd=_round_mode(), + part=pto.VcvtPartMode.EVEN, + ) + vec_i16 = pto.vcvt( + vec_i32, + pto.i16, + full_mask_b32, + sat=pto.VcvtSatMode.NOSAT, + part=pto.VcvtPartMode.EVEN, + ) + pto.vsts( + vec_i16, + dst_ptr, + offset, + store_mask, + dist=pto.VStoreDist.PK_B32, + ) + + _emit_tcvt_1d(dst, pto.f32, pto.i32, 1, emit_chunk) + + +template_tcvt_f16_to_i16_1d = _register_tcvt_1d( + name="template_tcvt_f16_to_i16", + dtypes=("f16", "i16"), + renderer=_render_tcvt_f16_to_i16_1d, +) template_tcvt_bf16_to_f16 = _register_tcvt( name="template_tcvt_bf16_to_f16", @@ -274,16 +500,32 @@ def template_tcvt_f16_to_i16(src: pto.Tile, dst: pto.Tile): def template_tcvt_f32_to_f32(src: pto.Tile, dst: pto.Tile): valid_rows, valid_cols = dst.valid_shape lanes_f32 = pto.elements_per_vreg(src.dtype) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols - col_loop = pto.for_(0, valid_cols, step=lanes_f32).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_f32): mask, remained = pto.make_mask(src.dtype, remained) vec = pto.vlds(src[row, col:]) converted = pto.vtrc(vec, mask, rnd=_round_mode()) pto.vsts(converted, dst[row, col:], mask) - col_loop.update(remained=remained) + + +def _render_tcvt_f32_to_f32_1d(src: pto.Tile, dst: pto.Tile): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + def emit_chunk(offset, mask): + vec = pto.vlds(src_ptr, offset) + converted = pto.vtrc(vec, mask, rnd=_round_mode()) + pto.vsts(converted, dst_ptr, offset, mask) + + _emit_tcvt_1d(dst, src.dtype, src.dtype, 1, emit_chunk) + + +template_tcvt_f32_to_f32_1d = _register_tcvt_1d( + name="template_tcvt_f32_to_f32", + dtypes=("f32", "f32"), + renderer=_render_tcvt_f32_to_f32_1d, +) @tilelib.tile_template( @@ -304,11 +546,9 @@ def template_tcvt_f32_to_i16(src: pto.Tile, dst: pto.Tile): valid_rows, valid_cols = dst.valid_shape lanes_f32 = pto.elements_per_vreg(src.dtype) full_mask = pto.make_mask(src.dtype, pto.PAT.ALL) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols - col_loop = pto.for_(0, valid_cols, step=lanes_f32).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_f32): store_mask, remained = pto.make_mask(src.dtype, remained) vec_f32 = pto.vlds(src[row, col:]) vec_i32 = pto.vcvt( @@ -326,7 +566,45 @@ def template_tcvt_f32_to_i16(src: pto.Tile, dst: pto.Tile): part=pto.VcvtPartMode.EVEN, ) pto.vsts(vec_i16, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B32) - col_loop.update(remained=remained) + + +def _render_tcvt_f32_to_i16_1d(src: pto.Tile, dst: pto.Tile): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + full_mask = pto.make_mask(src.dtype, pto.PAT.ALL) + + def emit_chunk(offset, store_mask): + vec_f32 = pto.vlds(src_ptr, offset) + vec_i32 = pto.vcvt( + vec_f32, + pto.i32, + full_mask, + rnd=_round_mode(), + sat=pto.VcvtSatMode.NOSAT, + ) + vec_i16 = pto.vcvt( + vec_i32, + pto.i16, + full_mask, + sat=pto.VcvtSatMode.NOSAT, + part=pto.VcvtPartMode.EVEN, + ) + pto.vsts( + vec_i16, + dst_ptr, + offset, + store_mask, + dist=pto.VStoreDist.PK_B32, + ) + + _emit_tcvt_1d(dst, src.dtype, src.dtype, 1, emit_chunk) + + +template_tcvt_f32_to_i16_1d = _register_tcvt_1d( + name="template_tcvt_f32_to_i16", + dtypes=("f32", "i16"), + renderer=_render_tcvt_f32_to_i16_1d, +) @tilelib.tile_template( @@ -347,11 +625,9 @@ def template_tcvt_f32_to_i64(src: pto.Tile, dst: pto.Tile): valid_rows, valid_cols = dst.valid_shape lanes_i64 = pto.elements_per_vreg(dst.dtype) full_mask = pto.make_mask(src.dtype, pto.PAT.ALL) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols * 2 - col_loop = pto.for_(0, valid_cols, step=lanes_i64).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_i64): store_mask, remained = pto.make_mask(pto.i32, remained) vec = pto.vlds(src[row, col:], dist="UNPK_B32") converted = pto.vcvt( @@ -363,7 +639,39 @@ def template_tcvt_f32_to_i64(src: pto.Tile, dst: pto.Tile): part=pto.VcvtPartMode.EVEN, ) pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.NORM_B32) - col_loop.update(remained=remained) + + +def _render_tcvt_f32_to_i64_1d(src: pto.Tile, dst: pto.Tile): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + full_mask = pto.make_mask(src.dtype, pto.PAT.ALL) + + def emit_chunk(offset, store_mask): + vec = pto.vlds(src_ptr, offset, dist="UNPK_B32") + converted = pto.vcvt( + vec, + pto.i64, + full_mask, + rnd=_round_mode(), + sat=pto.VcvtSatMode.SAT, + part=pto.VcvtPartMode.EVEN, + ) + pto.vsts( + converted, + dst_ptr, + offset, + store_mask, + dist=pto.VStoreDist.NORM_B32, + ) + + _emit_tcvt_1d(dst, dst.dtype, pto.i32, 2, emit_chunk) + + +template_tcvt_f32_to_i64_1d = _register_tcvt_1d( + name="template_tcvt_f32_to_i64", + dtypes=("f32", "i64"), + renderer=_render_tcvt_f32_to_i64_1d, +) template_tcvt_f16_to_i32 = _register_tcvt( name="template_tcvt_f16_to_i32", @@ -404,11 +712,9 @@ def template_tcvt_f16_to_ui8(src: pto.Tile, dst: pto.Tile): valid_rows, valid_cols = dst.valid_shape lanes_f16 = pto.elements_per_vreg(src.dtype) full_mask = pto.make_mask(src.dtype, pto.PAT.ALL) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols - col_loop = pto.for_(0, valid_cols, step=lanes_f16).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_f16): store_mask, remained = pto.make_mask(src.dtype, remained) vec = pto.vlds(src[row, col:]) converted = pto.vcvt( @@ -420,7 +726,39 @@ def template_tcvt_f16_to_ui8(src: pto.Tile, dst: pto.Tile): part=pto.VcvtPartMode.EVEN, ) pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B16) - col_loop.update(remained=remained) + + +def _render_tcvt_f16_to_ui8_1d(src: pto.Tile, dst: pto.Tile): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + full_mask = pto.make_mask(src.dtype, pto.PAT.ALL) + + def emit_chunk(offset, store_mask): + vec = pto.vlds(src_ptr, offset) + converted = pto.vcvt( + vec, + pto.ui8, + full_mask, + rnd=_round_mode(), + sat=pto.VcvtSatMode.NOSAT, + part=pto.VcvtPartMode.EVEN, + ) + pto.vsts( + converted, + dst_ptr, + offset, + store_mask, + dist=pto.VStoreDist.PK_B16, + ) + + _emit_tcvt_1d(dst, src.dtype, src.dtype, 1, emit_chunk) + + +template_tcvt_f16_to_ui8_1d = _register_tcvt_1d( + name="template_tcvt_f16_to_ui8", + dtypes=("f16", "ui8"), + renderer=_render_tcvt_f16_to_ui8_1d, +) @tilelib.tile_template( @@ -441,11 +779,9 @@ def template_tcvt_f16_to_si8(src: pto.Tile, dst: pto.Tile): valid_rows, valid_cols = dst.valid_shape lanes_f16 = pto.elements_per_vreg(src.dtype) pg = pto.make_mask(src.dtype, pto.PAT.ALL) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols - col_loop = pto.for_(0, valid_cols, step=lanes_f16).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_f16): full_mask, _ = pto.make_mask(src.dtype, lanes_f16) store_mask, remained = pto.make_mask(src.dtype, remained) vec_f16 = pto.vlds(src[row, col:]) @@ -473,7 +809,58 @@ def template_tcvt_f16_to_si8(src: pto.Tile, dst: pto.Tile): part=pto.VcvtPartMode.EVEN, ) pto.vsts(vec_si8, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B16) - col_loop.update(remained=remained) + + +def _render_tcvt_f16_to_si8_1d(src: pto.Tile, dst: pto.Tile): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + pg = pto.make_mask(src.dtype, pto.PAT.ALL) + full_mask, _ = pto.make_mask( + src.dtype, + pto.elements_per_vreg(src.dtype), + ) + + def emit_chunk(offset, store_mask): + vec_f16 = pto.vlds(src_ptr, offset) + vec_i16 = pto.vcvt( + vec_f16, + pto.i16, + full_mask, + rnd=_round_mode(), + sat=pto.VcvtSatMode.NOSAT, + ) + v_mask = pto.vdup(pto.i16(255), pg) + vec_i16_and = pto.vand(vec_i16, v_mask, store_mask) + vec_f16_temp = pto.vcvt( + vec_i16_and, + pto.f16, + full_mask, + rnd=_round_mode(), + ) + vec_si8 = pto.vcvt( + vec_f16_temp, + pto.si8, + full_mask, + rnd=_round_mode(), + sat=pto.VcvtSatMode.NOSAT, + part=pto.VcvtPartMode.EVEN, + ) + pto.vsts( + vec_si8, + dst_ptr, + offset, + store_mask, + dist=pto.VStoreDist.PK_B16, + ) + + _emit_tcvt_1d(dst, src.dtype, src.dtype, 1, emit_chunk) + + +template_tcvt_f16_to_si8_1d = _register_tcvt_1d( + name="template_tcvt_f16_to_si8", + dtypes=("f16", "si8"), + renderer=_render_tcvt_f16_to_si8_1d, +) template_tcvt_bf16_to_f32 = _register_tcvt( @@ -637,15 +1024,10 @@ def template_tcvt_si8_to_i32(src: pto.Tile, dst: pto.Tile): v_zero = pto.vbitcast(pto.vdup(pto.i8(0), b8_mask), pto.ui8) lanes_i16 = pto.elements_per_vreg(pto.i16) lanes_i32 = pto.elements_per_vreg(pto.i32) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols next_remained = valid_cols - lanes_i32 - col_loop = pto.for_(0, valid_cols, step=lanes_i16).carry( - remained=remained, - next_remained=next_remained, - ) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_i16): mask_b16_cur, remained = pto.make_mask(pto.i16, remained) mask_b16_next, next_remained = pto.make_mask(pto.i16, next_remained) mask_b32_cur = pto.punpack(mask_b16_cur, pto.PredicatePart.LOWER, to_type=pto.mask_b32) @@ -659,9 +1041,77 @@ def template_tcvt_si8_to_i32(src: pto.Tile, dst: pto.Tile): output_1 = pto.vcvt(vec_si8_2, pto.i32, b8_mask, part=pto.VcvtPartMode.P0) pto.vsts(output_0, dst[row, col:], mask_b32_cur, dist=pto.VStoreDist.NORM_B32) pto.vsts(output_1, dst[row, col + lanes_i32:], mask_b32_next, dist=pto.VStoreDist.NORM_B32) - col_loop.update(remained=remained, next_remained=next_remained) +@rewrite_jit_function +def _render_tcvt_si8_to_i32_1d(src: pto.Tile, dst: pto.Tile): + valid_rows, valid_cols = dst.valid_shape + total_elements = valid_rows * valid_cols + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + b8_mask = pto.make_mask(pto.ui8, pto.PAT.ALL) + v_zero = pto.vbitcast(pto.vdup(pto.i8(0), b8_mask), pto.ui8) + lanes_i16 = pto.elements_per_vreg(pto.i16) + lanes_i32 = pto.elements_per_vreg(pto.i32) + remained = total_elements + next_remained = total_elements - lanes_i32 + for offset in range(0, total_elements, lanes_i16): + mask_b16_cur, remained = pto.make_mask(pto.i16, remained) + mask_b16_next, next_remained = pto.make_mask( + pto.i16, + next_remained, + ) + mask_b32_cur = pto.punpack( + mask_b16_cur, + pto.PredicatePart.LOWER, + to_type=pto.mask_b32, + ) + mask_b32_next = pto.punpack( + mask_b16_next, + pto.PredicatePart.LOWER, + to_type=pto.mask_b32, + ) + vec_si8_0 = pto.vlds(src_ptr, offset, dist="UNPK_B8") + vec_ui8_0 = pto.vbitcast(vec_si8_0, pto.ui8) + vec_ui8_1, vec_ui8_2 = pto.vintlv(vec_ui8_0, v_zero) + vec_si8_1 = pto.vbitcast(vec_ui8_1, pto.si8) + vec_si8_2 = pto.vbitcast(vec_ui8_2, pto.si8) + output_0 = pto.vcvt( + vec_si8_1, + pto.i32, + b8_mask, + part=pto.VcvtPartMode.P0, + ) + output_1 = pto.vcvt( + vec_si8_2, + pto.i32, + b8_mask, + part=pto.VcvtPartMode.P0, + ) + pto.vsts( + output_0, + dst_ptr, + offset, + mask_b32_cur, + dist=pto.VStoreDist.NORM_B32, + ) + pto.vsts( + output_1, + dst_ptr, + offset + lanes_i32, + mask_b32_next, + dist=pto.VStoreDist.NORM_B32, + ) + + +template_tcvt_si8_to_i32_1d = _register_tcvt_1d( + name="template_tcvt_si8_to_i32", + dtypes=("si8", "i32"), + renderer=_render_tcvt_si8_to_i32_1d, +) + + +@rewrite_jit_function def _render_32_to_ui8(src: pto.Tile, dst: pto.Tile): valid_rows, valid_cols = dst.valid_shape full_mask = pto.make_mask(src.dtype, pto.PAT.ALL) @@ -672,12 +1122,13 @@ def _render_32_to_ui8(src: pto.Tile, dst: pto.Tile): v_idx_i16 = pto.vbitcast(v_idx, pto.i16) v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols - col_loop = pto.for_(0, valid_cols, step=lanes).carry(remained=remained) - with col_loop: - col = col_loop.iv - store_mask, remained = pto.make_mask(pto.ui8, remained) + for col in range(0, valid_cols, lanes): + # One source register contributes only 64 compacted bytes. Building + # this directly as a b8 mask would consume up to 256 elements. + iteration_mask, remained = pto.make_mask(src.dtype, remained) + store_mask = pto.pbitcast(iteration_mask, pto.mask_b8) vec = pto.vlds(src[row, col:]) converted = pto.vcvt( vec, @@ -689,7 +1140,41 @@ def _render_32_to_ui8(src: pto.Tile, dst: pto.Tile): result = pto.vselr(converted, v_idx_ui8) pto.mem_bar(pto.BarrierType.VST_VST) pto.vsts(result, dst[row, col:], store_mask, dist=pto.VStoreDist.NORM_B8) - col_loop.update(remained=remained) + + +def _render_32_to_ui8_1d(src: pto.Tile, dst: pto.Tile): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + full_mask = pto.make_mask(src.dtype, pto.PAT.ALL) + idx_mask_b8 = pto.pset_b8(pto.PAT.ALL) + idx_mask_b16 = pto.pbitcast(idx_mask_b8, pto.mask_b16) + v_idx = pto.vci(pto.i8(0), "ASC") + v_idx_i16 = pto.vbitcast(v_idx, pto.i16) + v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) + v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) + + def emit_chunk(offset, iteration_mask): + # Preserve the 64-element b32 loop granularity for the b8 store. + store_mask = pto.pbitcast(iteration_mask, pto.mask_b8) + vec = pto.vlds(src_ptr, offset) + converted = pto.vcvt( + vec, + pto.ui8, + full_mask, + sat=pto.VcvtSatMode.SAT, + part=pto.VcvtPartMode.P0, + ) + result = pto.vselr(converted, v_idx_ui8) + pto.mem_bar(pto.BarrierType.VST_VST) + pto.vsts( + result, + dst_ptr, + offset, + store_mask, + dist=pto.VStoreDist.NORM_B8, + ) + + _emit_tcvt_1d(dst, src.dtype, src.dtype, 1, emit_chunk) @tilelib.tile_template( @@ -728,6 +1213,19 @@ def template_tcvt_ui32_to_ui8(src: pto.Tile, dst: pto.Tile): _render_32_to_ui8(src, dst) +template_tcvt_i32_to_ui8_1d = _register_tcvt_1d( + name="template_tcvt_i32_to_ui8", + dtypes=("i32", "ui8"), + renderer=_render_32_to_ui8_1d, +) + +template_tcvt_ui32_to_ui8_1d = _register_tcvt_1d( + name="template_tcvt_ui32_to_ui8", + dtypes=("ui32", "ui8"), + renderer=_render_32_to_ui8_1d, +) + + @tilelib.tile_template( op="pto.tcvt", target="a5", @@ -746,11 +1244,9 @@ def template_tcvt_i16_to_ui8(src: pto.Tile, dst: pto.Tile): valid_rows, valid_cols = dst.valid_shape lanes_i16 = pto.elements_per_vreg(src.dtype) full_mask = pto.make_mask(src.dtype, pto.PAT.ALL) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols - col_loop = pto.for_(0, valid_cols, step=lanes_i16).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_i16): store_mask, remained = pto.make_mask(src.dtype, remained) vec = pto.vlds(src[row, col:]) converted = pto.vcvt( @@ -761,7 +1257,38 @@ def template_tcvt_i16_to_ui8(src: pto.Tile, dst: pto.Tile): part=pto.VcvtPartMode.EVEN, ) pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B16) - col_loop.update(remained=remained) + + +def _render_tcvt_i16_to_ui8_1d(src: pto.Tile, dst: pto.Tile): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + full_mask = pto.make_mask(src.dtype, pto.PAT.ALL) + + def emit_chunk(offset, store_mask): + vec = pto.vlds(src_ptr, offset) + converted = pto.vcvt( + vec, + pto.ui8, + full_mask, + sat=pto.VcvtSatMode.SAT, + part=pto.VcvtPartMode.EVEN, + ) + pto.vsts( + converted, + dst_ptr, + offset, + store_mask, + dist=pto.VStoreDist.PK_B16, + ) + + _emit_tcvt_1d(dst, src.dtype, src.dtype, 1, emit_chunk) + + +template_tcvt_i16_to_ui8_1d = _register_tcvt_1d( + name="template_tcvt_i16_to_ui8", + dtypes=("i16", "ui8"), + renderer=_render_tcvt_i16_to_ui8_1d, +) @tilelib.tile_template( @@ -782,11 +1309,9 @@ def template_tcvt_i32_to_i64(src: pto.Tile, dst: pto.Tile): valid_rows, valid_cols = dst.valid_shape lanes_i64 = pto.elements_per_vreg(dst.dtype) full_mask = pto.make_mask(src.dtype, pto.PAT.ALL) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols * 2 - col_loop = pto.for_(0, valid_cols, step=lanes_i64).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_i64): store_mask, remained = pto.make_mask(pto.i32, remained) vec = pto.vlds(src[row, col:], dist="UNPK_B32") converted = pto.vcvt( @@ -796,28 +1321,108 @@ def template_tcvt_i32_to_i64(src: pto.Tile, dst: pto.Tile): part=pto.VcvtPartMode.EVEN, ) pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.NORM_B32) - col_loop.update(remained=remained) +def _render_tcvt_i32_to_i64_1d(src: pto.Tile, dst: pto.Tile): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + full_mask = pto.make_mask(src.dtype, pto.PAT.ALL) + + def emit_chunk(offset, store_mask): + vec = pto.vlds(src_ptr, offset, dist="UNPK_B32") + converted = pto.vcvt( + vec, + pto.i64, + full_mask, + part=pto.VcvtPartMode.EVEN, + ) + pto.vsts( + converted, + dst_ptr, + offset, + store_mask, + dist=pto.VStoreDist.NORM_B32, + ) + + _emit_tcvt_1d(dst, dst.dtype, pto.i32, 2, emit_chunk) + + +template_tcvt_i32_to_i64_1d = _register_tcvt_1d( + name="template_tcvt_i32_to_i64", + dtypes=("i32", "i64"), + renderer=_render_tcvt_i32_to_i64_1d, +) + + +@rewrite_jit_function def _render_i64_to_32(src: pto.Tile, dst: pto.Tile, *, use_rounding: bool): valid_rows, valid_cols = dst.valid_shape lanes_i64 = pto.elements_per_vreg(src.dtype) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols * 2 full_mask, _ = pto.make_mask(pto.i32, remained) - col_loop = pto.for_(0, valid_cols, step=lanes_i64).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_i64): store_mask, remained = pto.make_mask(dst.dtype, remained) vec = pto.vlds(src[row, col:]) - kwargs = {"part": pto.VcvtPartMode.EVEN} if use_rounding: - kwargs["rnd"] = _round_mode() + converted = pto.vcvt( + vec, + dst.dtype, + full_mask, + rnd=_round_mode(), + part=pto.VcvtPartMode.EVEN, + ) else: - kwargs["sat"] = pto.VcvtSatMode.NOSAT - converted = pto.vcvt(vec, dst.dtype, full_mask, **kwargs) + converted = pto.vcvt( + vec, + dst.dtype, + full_mask, + sat=pto.VcvtSatMode.NOSAT, + part=pto.VcvtPartMode.EVEN, + ) pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B64) - col_loop.update(remained=remained) + + +def _render_i64_to_32_1d(src: pto.Tile, dst: pto.Tile, *, use_rounding: bool): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + full_mask = pto.make_mask(pto.i32, pto.PAT.ALL) + + def emit_chunk(offset, store_mask): + vec = pto.vlds(src_ptr, offset) + if use_rounding: + converted = pto.vcvt( + vec, + dst.dtype, + full_mask, + rnd=_round_mode(), + part=pto.VcvtPartMode.EVEN, + ) + else: + converted = pto.vcvt( + vec, + dst.dtype, + full_mask, + sat=pto.VcvtSatMode.NOSAT, + part=pto.VcvtPartMode.EVEN, + ) + pto.vsts( + converted, + dst_ptr, + offset, + store_mask, + dist=pto.VStoreDist.PK_B64, + ) + + _emit_tcvt_1d(dst, src.dtype, dst.dtype, 2, emit_chunk) + + +def _render_i64_to_f32_1d(src: pto.Tile, dst: pto.Tile): + _render_i64_to_32_1d(src, dst, use_rounding=True) + + +def _render_i64_to_i32_1d(src: pto.Tile, dst: pto.Tile): + _render_i64_to_32_1d(src, dst, use_rounding=False) @tilelib.tile_template( @@ -856,6 +1461,19 @@ def template_tcvt_i64_to_i32(src: pto.Tile, dst: pto.Tile): _render_i64_to_32(src, dst, use_rounding=False) +template_tcvt_i64_to_f32_1d = _register_tcvt_1d( + name="template_tcvt_i64_to_f32", + dtypes=("i64", "f32"), + renderer=_render_i64_to_f32_1d, +) + +template_tcvt_i64_to_i32_1d = _register_tcvt_1d( + name="template_tcvt_i64_to_i32", + dtypes=("i64", "i32"), + renderer=_render_i64_to_i32_1d, +) + + @tilelib.tile_template( op="pto.tcvt", target="a5", @@ -881,11 +1499,9 @@ def template_tcvt_f32_to_fp8(src: pto.Tile, dst: pto.Tile): v_idx_i16 = pto.vbitcast(v_idx, pto.i16) v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols - col_loop = pto.for_(0, valid_cols, step=lanes_f32).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_f32): dst_mask, remained = pto.make_mask(dst_dtype, remained) vec = pto.vlds(src[row, col:]) converted = pto.vcvt( @@ -899,7 +1515,49 @@ def template_tcvt_f32_to_fp8(src: pto.Tile, dst: pto.Tile): result = _vselr_low_precision(converted, v_idx_ui8) pto.mem_bar(pto.BarrierType.VST_VST) pto.vsts(result, dst[row, col:], dst_mask, dist=pto.VStoreDist.NORM_B8) - col_loop.update(remained=remained) + + +def _render_tcvt_f32_to_fp8_1d(src: pto.Tile, dst: pto.Tile): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + dst_dtype = dst.dtype + src_mask = pto.make_mask(src.dtype, pto.PAT.ALL) + idx_mask_b8 = pto.pset_b8(pto.PAT.ALL) + idx_mask_b16 = pto.pbitcast(idx_mask_b8, pto.mask_b16) + v_idx = pto.vci(pto.i8(0), "ASC") + v_idx_i16 = pto.vbitcast(v_idx, pto.i16) + v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) + v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) + + def emit_chunk(offset, dst_mask): + vec = pto.vlds(src_ptr, offset) + converted = pto.vcvt( + vec, + dst_dtype, + src_mask, + rnd=_round_mode(), + sat=pto.VcvtSatMode.SAT, + part=pto.VcvtPartMode.P0, + ) + result = _vselr_low_precision(converted, v_idx_ui8) + pto.mem_bar(pto.BarrierType.VST_VST) + pto.vsts( + result, + dst_ptr, + offset, + dst_mask, + dist=pto.VStoreDist.NORM_B8, + ) + + _emit_tcvt_1d(dst, src.dtype, dst_dtype, 1, emit_chunk) + + +template_tcvt_f32_to_fp8_1d = _register_tcvt_1d( + name="template_tcvt_f32_to_fp8", + dtypes=(("f32", "f8e4m3"), ("f32", "f8e5m2")), + renderer=_render_tcvt_f32_to_fp8_1d, + tags=("low_precision",), +) @tilelib.tile_template( @@ -927,11 +1585,9 @@ def template_tcvt_f32_to_hif8(src: pto.Tile, dst: pto.Tile): v_idx_i16 = pto.vbitcast(v_idx, pto.i16) v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols - col_loop = pto.for_(0, valid_cols, step=lanes_f32).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_f32): dst_mask, remained = pto.make_mask(dst_dtype, remained) vec = pto.vlds(src[row, col:]) converted = pto.vcvt( @@ -945,7 +1601,49 @@ def template_tcvt_f32_to_hif8(src: pto.Tile, dst: pto.Tile): result = _vselr_low_precision(converted, v_idx_ui8) pto.mem_bar(pto.BarrierType.VST_VST) pto.vsts(result, dst[row, col:], dst_mask, dist=pto.VStoreDist.NORM_B8) - col_loop.update(remained=remained) + + +def _render_tcvt_f32_to_hif8_1d(src: pto.Tile, dst: pto.Tile): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + dst_dtype = dst.dtype + src_mask = pto.make_mask(src.dtype, pto.PAT.ALL) + idx_mask_b8 = pto.pset_b8(pto.PAT.ALL) + idx_mask_b16 = pto.pbitcast(idx_mask_b8, pto.mask_b16) + v_idx = pto.vci(pto.i8(0), "ASC") + v_idx_i16 = pto.vbitcast(v_idx, pto.i16) + v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) + v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) + + def emit_chunk(offset, dst_mask): + vec = pto.vlds(src_ptr, offset) + converted = pto.vcvt( + vec, + dst_dtype, + src_mask, + rnd=pto.VcvtRoundMode.A, + sat=pto.VcvtSatMode.NOSAT, + part=pto.VcvtPartMode.P0, + ) + result = _vselr_low_precision(converted, v_idx_ui8) + pto.mem_bar(pto.BarrierType.VST_VST) + pto.vsts( + result, + dst_ptr, + offset, + dst_mask, + dist=pto.VStoreDist.NORM_B8, + ) + + _emit_tcvt_1d(dst, src.dtype, dst_dtype, 1, emit_chunk) + + +template_tcvt_f32_to_hif8_1d = _register_tcvt_1d( + name="template_tcvt_f32_to_hif8", + dtypes=("f32", "hif8"), + renderer=_render_tcvt_f32_to_hif8_1d, + tags=("low_precision",), +) @tilelib.tile_template( @@ -967,11 +1665,9 @@ def template_tcvt_f16_to_hif8(src: pto.Tile, dst: pto.Tile): dst_dtype = dst.dtype lanes_f16 = pto.elements_per_vreg(src.dtype) src_mask = pto.make_mask(src.dtype, pto.PAT.ALL) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols - col_loop = pto.for_(0, valid_cols, step=lanes_f16).carry(remained=remained) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, lanes_f16): dst_mask, remained = pto.make_mask(src.dtype, remained) vec = pto.vlds(src[row, col:]) converted = pto.vcvt( @@ -983,7 +1679,41 @@ def template_tcvt_f16_to_hif8(src: pto.Tile, dst: pto.Tile): part=pto.VcvtPartMode.EVEN, ) pto.vsts(converted, dst[row, col:], dst_mask, dist=pto.VStoreDist.PK_B16) - col_loop.update(remained=remained) + + +def _render_tcvt_f16_to_hif8_1d(src: pto.Tile, dst: pto.Tile): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + dst_dtype = dst.dtype + src_mask = pto.make_mask(src.dtype, pto.PAT.ALL) + + def emit_chunk(offset, dst_mask): + vec = pto.vlds(src_ptr, offset) + converted = pto.vcvt( + vec, + dst_dtype, + src_mask, + rnd=pto.VcvtRoundMode.A, + sat=pto.VcvtSatMode.NOSAT, + part=pto.VcvtPartMode.EVEN, + ) + pto.vsts( + converted, + dst_ptr, + offset, + dst_mask, + dist=pto.VStoreDist.PK_B16, + ) + + _emit_tcvt_1d(dst, src.dtype, src.dtype, 1, emit_chunk) + + +template_tcvt_f16_to_hif8_1d = _register_tcvt_1d( + name="template_tcvt_f16_to_hif8", + dtypes=("f16", "hif8"), + renderer=_render_tcvt_f16_to_hif8_1d, + tags=("low_precision",), +) @tilelib.tile_template( @@ -1011,15 +1741,10 @@ def template_tcvt_bf16_to_fp4(src: pto.Tile, dst: pto.Tile): v_idx_i16 = pto.vbitcast(v_idx, pto.i16) v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) - with pto.for_(0, valid_rows, step=1) as row: + for row in range(0, valid_rows, 1): remained = valid_cols src_remained = valid_cols * 2 - col_loop = pto.for_(0, valid_cols, step=dst_chunk_cols).carry( - remained=remained, - src_remained=src_remained, - ) - with col_loop: - col = col_loop.iv + for col in range(0, valid_cols, dst_chunk_cols): dst_mask, remained = pto.make_mask(dst_dtype, remained) src_mask, src_remained = pto.make_mask(src.dtype, src_remained) vec = pto.vlds(src[row, col * 2:]) @@ -1033,4 +1758,94 @@ def template_tcvt_bf16_to_fp4(src: pto.Tile, dst: pto.Tile): result = _vselr_low_precision(converted, v_idx_ui8) pto.mem_bar(pto.BarrierType.VST_VST) pto.vsts(result, dst[row, col:], dst_mask, dist=pto.VStoreDist.NORM_B8) - col_loop.update(remained=remained, src_remained=src_remained) + + +@rewrite_jit_function +def _render_tcvt_bf16_to_fp4_1d(src: pto.Tile, dst: pto.Tile): + valid_rows, valid_cols = dst.valid_shape + total_destination_elements = valid_rows * valid_cols + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + dst_dtype = dst.dtype + lanes_bf16 = pto.elements_per_vreg(src.dtype) + dst_chunk_cols = lanes_bf16 // 2 + idx_mask_b8 = pto.pset_b8(pto.PAT.ALL) + idx_mask_b16 = pto.pbitcast(idx_mask_b8, pto.mask_b16) + v_idx = pto.vci(pto.i8(0), "ASC") + v_idx_i16 = pto.vbitcast(v_idx, pto.i16) + v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) + v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) + remained = total_destination_elements + src_remained = total_destination_elements * 2 + for offset in range(0, total_destination_elements, dst_chunk_cols): + dst_mask, remained = pto.make_mask(dst_dtype, remained) + src_mask, src_remained = pto.make_mask(src.dtype, src_remained) + vec = pto.vlds(src_ptr, offset * 2) + converted = pto.vcvt( + vec, + dst_dtype, + src_mask, + rnd=_round_mode(), + part=pto.VcvtPartMode.P0, + ) + result = _vselr_low_precision(converted, v_idx_ui8) + pto.mem_bar(pto.BarrierType.VST_VST) + pto.vsts( + result, + dst_ptr, + offset, + dst_mask, + dist=pto.VStoreDist.NORM_B8, + ) + + +template_tcvt_bf16_to_fp4_1d = _register_tcvt_1d( + name="template_tcvt_bf16_to_fp4", + dtypes=( + ("bf16", "f4e1m2x2"), + ("bf16", "f4e2m1x2"), + ), + renderer=_render_tcvt_bf16_to_fp4_1d, + source_elements_per_destination=2, + tags=("low_precision",), +) + + +def _register_deferred_tcvt_1d(): + """Assign stable 1D IDs after every 2D fallback is registered.""" + + registry = tilelib.default_registry() + fallbacks = { + candidate.name: candidate + for candidate in registry.lookup("pto.tcvt", "a5") + if candidate.metadata.loop_depth == 2 + } + next_candidate_id = max( + candidate.metadata.id for candidate in fallbacks.values() + ) + 1 + replacements = {} + for candidate in sorted( + _PENDING_TCVT_1D, + key=lambda item: fallbacks[ + item.name.removesuffix("_1d") + ].metadata.id, + ): + assigned = replace( + candidate, + metadata=replace( + candidate.metadata, + id=next_candidate_id, + ), + ) + registry.register(assigned) + replacements[id(candidate)] = assigned + next_candidate_id += 1 + + for global_name, value in tuple(globals().items()): + assigned = replacements.get(id(value)) + if assigned is not None: + globals()[global_name] = assigned + _PENDING_TCVT_1D.clear() + + +_register_deferred_tcvt_1d() diff --git a/lib/TileOps/a5/tdiv.py b/lib/TileOps/a5/tdiv.py index 3d1f6c8378..6e523d606e 100644 --- a/lib/TileOps/a5/tdiv.py +++ b/lib/TileOps/a5/tdiv.py @@ -10,41 +10,69 @@ from ptodsl import pto import ptodsl.tilelib as tilelib +from ._elementwise import emit_binary_1d, emit_binary_2d, traversal_metadata from .div_hp import _div_ieee754_f32_impl, _div_ieee754_f16_impl -@tilelib.tile_template( - op="pto.tdiv", - target="a5", - name="template_tdiv", - dtypes=[("f16", "f16", "f16"), ("f32", "f32", "f32")], - iteration_axis="none", - op_engine="vector", - op_class="elementwise", - layouts=["row_major"], - memory_spaces=["ub"], - priority=0, - id=0, - loop_depth=2, - is_post_update=False, -) -def template_tdiv(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): +_DTYPES = [("f16", "f16", "f16"), ("f32", "f32", "f32")] + + +def _emit_tdiv(src0, src1, dst, traversal): + """Emit the operation-specific default or high-precision division.""" + dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - lanes = pto.elements_per_vreg(dtype) precision_type = pto.get_op_attr("precisionType", "default") - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - if precision_type == "high_precision": - if str(dtype) == "f32": - divided = _div_ieee754_f32_impl(lhs, rhs, mask) - else: - divided = _div_ieee754_f16_impl(lhs, rhs, mask) - else: - divided = pto.vdiv(lhs, rhs, mask) - pto.vsts(divided, dst[row, col:], mask) + def divide(lhs, rhs, mask): + if precision_type == "high_precision": + if str(dtype) == "f32": + return _div_ieee754_f32_impl(lhs, rhs, mask) + return _div_ieee754_f16_impl(lhs, rhs, mask) + return pto.vdiv(lhs, rhs, mask) + + if traversal == "1d": + emit_binary_1d(src0, src1, dst, divide) + else: + emit_binary_2d(src0, src1, dst, divide) + + +def _register_tdiv(*, name, traversal): + constraints = [] + loop_depth, priority, candidate_id = traversal_metadata(traversal) + if traversal == "1d": + constraints.append( + tilelib.require_elementwise_1d("src0", "src1", "dst") + ) + + @tilelib.tile_template( + op="pto.tdiv", + target="a5", + name=name, + dtypes=_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="elementwise", + layouts=["row_major"], + memory_spaces=["ub"], + constraints=constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, + is_post_update=False, + ) + def template(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): + _emit_tdiv(src0, src1, dst, traversal) + + return template + + +template_tdiv = _register_tdiv( + name="template_tdiv", + traversal="2d", +) + + +template_tdiv_1d = _register_tdiv( + name="template_tdiv_1d", + traversal="1d", +) diff --git a/lib/TileOps/a5/tdivs.py b/lib/TileOps/a5/tdivs.py index 933f2b5722..44e09643a2 100644 --- a/lib/TileOps/a5/tdivs.py +++ b/lib/TileOps/a5/tdivs.py @@ -10,7 +10,12 @@ from ptodsl import pto import ptodsl.tilelib as tilelib -from ._elementwise import _common_constraints +from ._elementwise import ( + _common_constraints, + emit_scalar_binary_1d, + emit_scalar_binary_2d, + traversal_metadata, +) from .div_hp import _div_ieee754_f32_impl, _div_ieee754_f16_impl @@ -36,62 +41,106 @@ def _div(lhs, rhs, dtype, mask, precision_type): return pto.vdiv(lhs, rhs, mask) -def _emit_tdivs_body(src, scalar, dst, *, scalar_lhs=False): +def _emit_tdivs_body(src, scalar, dst, traversal, *, scalar_lhs=False): dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - src_cols = src.shape[1] - dst_cols = dst.shape[1] - lanes = pto.elements_per_vreg(dtype) precision_type = pto.get_op_attr("precisionType", "default") - src_ptr = src.as_ptr() - dst_ptr = dst.as_ptr() - - with pto.for_(0, valid_rows, step=1) as row: - col_loop = pto.for_(0, valid_cols, step=lanes).carry(remained=valid_cols) - with col_loop: - col = col_loop.iv - mask, remained = pto.make_mask(dtype, col_loop.remained) - src_addr = pto.addptr(src_ptr, row * src_cols + col) - value = pto.vlds(src_addr, 0) - scalar_value = pto.vbr(scalar) - lhs, rhs = (scalar_value, value) if scalar_lhs else (value, scalar_value) - result = _div(lhs, rhs, dtype, mask, precision_type) - dst_addr = pto.addptr(dst_ptr, row * dst_cols + col) - pto.vsts(result, dst_addr, 0, mask) - col_loop.update(remained=remained) - - -@tilelib.tile_template( - op="pto.tdivs", - target="a5", + + def divide(lhs, rhs, mask): + return _div(lhs, rhs, dtype, mask, precision_type) + + emitter = ( + emit_scalar_binary_1d + if traversal == "1d" + else emit_scalar_binary_2d + ) + emitter( + src, + scalar, + dst, + divide, + broadcast_scalar=True, + scalar_lhs=scalar_lhs, + ) + + +def _register_tdivs(*, name, traversal, scalar_lhs=False): + constraints = _common_constraints("src", "dst") + constraints.append(_scalar_tile_tile if scalar_lhs else _tile_scalar_tile) + loop_depth, priority, candidate_id = traversal_metadata( + traversal, + fallback_candidate_id=1 if scalar_lhs else 0, + candidate_count=2, + ) + if traversal == "1d": + constraints.append(tilelib.require_elementwise_1d("src", "dst")) + + if scalar_lhs: + + @tilelib.tile_template( + op="pto.tdivs", + target="a5", + name=name, + dtypes=_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="elementwise", + constraints=constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, + is_post_update=False, + tags=("elementwise", "scalar"), + ) + def template(scalar, src: pto.Tile, dst: pto.Tile): + _emit_tdivs_body( + src, + scalar, + dst, + traversal, + scalar_lhs=True, + ) + + return template + + @tilelib.tile_template( + op="pto.tdivs", + target="a5", + name=name, + dtypes=_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="elementwise", + constraints=constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, + is_post_update=False, + tags=("elementwise", "scalar"), + ) + def template(src: pto.Tile, scalar, dst: pto.Tile): + _emit_tdivs_body(src, scalar, dst, traversal) + + return template + + +template_tdivs_tile_scalar = _register_tdivs( name="template_tdivs_tile_scalar", - dtypes=_DTYPES, - iteration_axis="none", - op_engine="vector", - op_class="elementwise", - constraints=_common_constraints("src", "dst") + [_tile_scalar_tile], - id=0, - loop_depth=2, - is_post_update=False, - tags=("elementwise", "scalar"), + traversal="2d", ) -def template_tdivs_tile_scalar(src: pto.Tile, scalar, dst: pto.Tile): - _emit_tdivs_body(src, scalar, dst) - -@tilelib.tile_template( - op="pto.tdivs", - target="a5", +template_tdivs_scalar_tile = _register_tdivs( name="template_tdivs_scalar_tile", - dtypes=_DTYPES, - iteration_axis="none", - op_engine="vector", - op_class="elementwise", - constraints=_common_constraints("src", "dst") + [_scalar_tile_tile], - id=1, - loop_depth=2, - is_post_update=False, - tags=("elementwise", "scalar"), + traversal="2d", + scalar_lhs=True, +) + +template_tdivs_tile_scalar_1d = _register_tdivs( + name="template_tdivs_tile_scalar_1d", + traversal="1d", +) + +template_tdivs_scalar_tile_1d = _register_tdivs( + name="template_tdivs_scalar_tile_1d", + traversal="1d", + scalar_lhs=True, ) -def template_tdivs_scalar_tile(scalar, src: pto.Tile, dst: pto.Tile): - _emit_tdivs_body(src, scalar, dst, scalar_lhs=True) diff --git a/lib/TileOps/a5/texp.py b/lib/TileOps/a5/texp.py index dc9b7c96ca..e81f900455 100644 --- a/lib/TileOps/a5/texp.py +++ b/lib/TileOps/a5/texp.py @@ -12,12 +12,24 @@ from ._elementwise import register_unary +_DTYPES = [ + ("f16", "f16"), + ("f32", "f32"), +] + + template_texp = register_unary( op="pto.texp", name="template_texp", vector_op=pto.vexp, - dtypes=[ - ("f16", "f16"), - ("f32", "f32"), - ], + dtypes=_DTYPES, +) + + +template_texp_1d = register_unary( + op="pto.texp", + name="template_texp_1d", + vector_op=pto.vexp, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/texpand.py b/lib/TileOps/a5/texpand.py index e405a6392b..ceb0423238 100644 --- a/lib/TileOps/a5/texpand.py +++ b/lib/TileOps/a5/texpand.py @@ -10,15 +10,25 @@ from ._elementwise import register_scalar_fill +_DTYPES = [ + ("i8", "i8"), + ("i16", "i16"), + ("i32", "i32"), + ("f16", "f16"), + ("bf16", "bf16"), + ("f32", "f32"), +] + + template_texpands = register_scalar_fill( op="pto.texpands", name="template_texpands", - dtypes=[ - ("i8", "i8"), - ("i16", "i16"), - ("i32", "i32"), - ("f16", "f16"), - ("bf16", "bf16"), - ("f32", "f32"), - ], + dtypes=_DTYPES, +) + +template_texpands_1d = register_scalar_fill( + op="pto.texpands", + name="template_texpands_1d", + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tfmod.py b/lib/TileOps/a5/tfmod.py index 32cc247592..9b1cce0e5d 100644 --- a/lib/TileOps/a5/tfmod.py +++ b/lib/TileOps/a5/tfmod.py @@ -16,3 +16,12 @@ dtypes=FMOD_DTYPES, round_mode="Z", ) + + +template_tfmod_1d = register_binary_remainder( + op="pto.tfmod", + name="template_tfmod_1d", + dtypes=FMOD_DTYPES, + round_mode="Z", + traversal="1d", +) diff --git a/lib/TileOps/a5/tfmods.py b/lib/TileOps/a5/tfmods.py index 7a4e4b6b87..19009b4cde 100644 --- a/lib/TileOps/a5/tfmods.py +++ b/lib/TileOps/a5/tfmods.py @@ -16,3 +16,11 @@ dtypes=FMODS_DTYPES, round_mode="Z", ) + +template_tfmods_1d = register_scalar_remainder( + op="pto.tfmods", + name="template_tfmods_1d", + dtypes=FMODS_DTYPES, + round_mode="Z", + traversal="1d", +) diff --git a/lib/TileOps/a5/tlog.py b/lib/TileOps/a5/tlog.py index a15220354b..a7134c3b7a 100644 --- a/lib/TileOps/a5/tlog.py +++ b/lib/TileOps/a5/tlog.py @@ -10,7 +10,19 @@ from ptodsl import pto import ptodsl.tilelib as tilelib -from ._elementwise import _common_constraints, register_unary +from ._elementwise import ( + _common_constraints, + emit_unary_1d, + emit_unary_2d, + register_unary, + traversal_metadata, +) + + +_DTYPES = [ + ("f16", "f16"), + ("f32", "f32"), +] def _is_default_precision(precisionType="default", **_): @@ -25,34 +37,26 @@ def _is_high_precision(precisionType="default", **_): op="pto.tlog", name="template_tlog", vector_op=pto.vln, - dtypes=[ - ("f16", "f16"), - ("f32", "f32"), - ], + dtypes=_DTYPES, constraints=[_is_default_precision], ) -@tilelib.tile_template( +template_tlog_1d = register_unary( op="pto.tlog", - target="a5", - name="template_tlog_high_precision", - dtypes=[ - ("f16", "f16"), - ("f32", "f32"), - ], - iteration_axis="none", - op_engine="vector", - op_class="elementwise", - constraints=_common_constraints("src", "dst") + [_is_high_precision], - id=1, - loop_depth=2, - is_post_update=False, - tags=("elementwise", "unary"), + name="template_tlog_1d", + vector_op=pto.vln, + dtypes=_DTYPES, + constraints=[_is_default_precision], + traversal="1d", + candidate_id=2, ) -def template_tlog_high_precision(src: pto.Tile, dst: pto.Tile): + + +def _emit_tlog_high_precision(src, dst, traversal): + """Emit the operation-specific high-precision logarithm computation.""" + dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape if str(dtype) == "f16": subnormal_threshold = pto.f16("0x03FF") mul_factor = pto.f16("0x6400") @@ -62,18 +66,67 @@ def template_tlog_high_precision(src: pto.Tile, dst: pto.Tile): mul_factor = pto.f32("0x4B000000") compensation = pto.f32(-15.9423851528787421) - lanes = pto.elements_per_vreg(dtype) - with pto.for_(0, valid_rows, step=1) as row: - col_loop = pto.for_(0, valid_cols, step=lanes).carry(remained=valid_cols) - with col_loop: - col = col_loop.iv - mask, remained = pto.make_mask(dtype, col_loop.remained) - vinput = pto.vlds(src[row, col:]) - cmp_mask = pto.vcmps(vinput, subnormal_threshold, mask, pto.CmpMode.LT) - scaled = pto.vmuls(vinput, mul_factor, mask) - selected_input = pto.vsel(scaled, vinput, cmp_mask) - log_result = pto.vln(selected_input, mask) - compensated = pto.vadds(log_result, compensation, mask) - result = pto.vsel(compensated, log_result, cmp_mask) - pto.vsts(result, dst[row, col:], mask) - col_loop.update(remained=remained) + def high_precision_log(value, mask): + cmp_mask = pto.vcmps( + value, + subnormal_threshold, + mask, + pto.CmpMode.LT, + ) + scaled = pto.vmuls(value, mul_factor, mask) + selected_input = pto.vsel(scaled, value, cmp_mask) + log_result = pto.vln(selected_input, mask) + compensated = pto.vadds(log_result, compensation, mask) + return pto.vsel(compensated, log_result, cmp_mask) + + if traversal == "1d": + emit_unary_1d(src, dst, high_precision_log) + else: + emit_unary_2d(src, dst, high_precision_log) + + +def _register_tlog_high_precision( + *, + name, + traversal, +): + constraints = _common_constraints("src", "dst") + [_is_high_precision] + loop_depth, priority, candidate_id = traversal_metadata( + traversal, + fallback_candidate_id=1, + candidate_count=2, + ) + if traversal == "1d": + constraints.append(tilelib.require_elementwise_1d("src", "dst")) + + @tilelib.tile_template( + op="pto.tlog", + target="a5", + name=name, + dtypes=_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="elementwise", + constraints=constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, + is_post_update=False, + tags=("elementwise", "unary"), + ) + def template(src: pto.Tile, dst: pto.Tile): + _emit_tlog_high_precision(src, dst, traversal) + + return template + + +template_tlog_high_precision = _register_tlog_high_precision( + name="template_tlog_high_precision", + traversal="2d", +) + + +template_tlog_high_precision_1d = _register_tlog_high_precision( + name="template_tlog_high_precision_1d", + traversal="1d", +) diff --git a/lib/TileOps/a5/tlrelu.py b/lib/TileOps/a5/tlrelu.py index 5efcfb1195..59b88b3c8c 100644 --- a/lib/TileOps/a5/tlrelu.py +++ b/lib/TileOps/a5/tlrelu.py @@ -11,6 +11,12 @@ import ptodsl.tilelib as tilelib from ptodsl._types import _resolve +from ._elementwise import ( + emit_scalar_binary_1d, + emit_scalar_binary_2d, + traversal_metadata, +) + def _ub_or_vec_row_major(operand_memory_spaces, operand_b_layouts, operand_s_layouts, **_): return ( @@ -20,43 +26,63 @@ def _ub_or_vec_row_major(operand_memory_spaces, operand_b_layouts, operand_s_lay ) -@tilelib.tile_template( - op="pto.tlrelu", - target="a5", - name="template_tlrelu", - dtypes=[ - ("f16", "f16", "f16"), - ("f16", "f32", "f16"), - ("f32", "f32", "f32"), - ], - iteration_axis="none", - op_engine="vector", - op_class="elementwise", - constraints=[ +_DTYPES = [ + ("f16", "f16", "f16"), + ("f16", "f32", "f16"), + ("f32", "f32", "f32"), +] + + +def _register_tlrelu(*, name, traversal): + constraints = [ _ub_or_vec_row_major, tilelib.require_same_valid_shape("src", "dst"), - ], - id=0, - loop_depth=2, - is_post_update=False, - tags=("elementwise", "scalar"), -) -def template_tlrelu(src: pto.Tile, slope, dst: pto.Tile): - dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - lanes = pto.elements_per_vreg(dtype) - slope_scalar = slope - if str(dtype) == "f16": - slope_scalar = scalar.coerce_scalar_to_type( - slope, - _resolve(pto.f16), - context="template_tlrelu(slope)", + ] + loop_depth, priority, candidate_id = traversal_metadata(traversal) + if traversal == "1d": + constraints.append(tilelib.require_elementwise_1d("src", "dst")) + + @tilelib.tile_template( + op="pto.tlrelu", + target="a5", + name=name, + dtypes=_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="elementwise", + constraints=constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, + is_post_update=False, + tags=("elementwise", "scalar"), + ) + def template(src: pto.Tile, slope, dst: pto.Tile): + dtype = dst.dtype + slope_scalar = slope + if str(dtype) == "f16": + slope_scalar = scalar.coerce_scalar_to_type( + slope, + _resolve(pto.f16), + context="template_tlrelu(slope)", + ) + + emitter = ( + emit_scalar_binary_1d + if traversal == "1d" + else emit_scalar_binary_2d ) + emitter(src, slope_scalar, dst, pto.vlrelu) + + return template - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - value = pto.vlds(src[row, col:]) - result = pto.vlrelu(value, slope_scalar, mask) - pto.vsts(result, dst[row, col:], mask) + +template_tlrelu = _register_tlrelu( + name="template_tlrelu", + traversal="2d", +) + +template_tlrelu_1d = _register_tlrelu( + name="template_tlrelu_1d", + traversal="1d", +) diff --git a/lib/TileOps/a5/tmax.py b/lib/TileOps/a5/tmax.py index c28f4e75f3..cc3958fea6 100644 --- a/lib/TileOps/a5/tmax.py +++ b/lib/TileOps/a5/tmax.py @@ -17,9 +17,21 @@ def _vmax(lhs, rhs, mask): return pto.vmax(lhs, rhs, mask) +_DTYPES = same_dtype_signatures(3) + + template_tmax = register_binary( op="pto.tmax", name="template_tmax", vector_op=_vmax, - dtypes=same_dtype_signatures(3), + dtypes=_DTYPES, +) + + +template_tmax_1d = register_binary( + op="pto.tmax", + name="template_tmax_1d", + vector_op=_vmax, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tmaxs.py b/lib/TileOps/a5/tmaxs.py index b8a5ef905a..ef79067ee9 100644 --- a/lib/TileOps/a5/tmaxs.py +++ b/lib/TileOps/a5/tmaxs.py @@ -13,9 +13,20 @@ from ._elementwise import register_scalar_binary +_DTYPES = same_dtype_signatures(3) + + template_tmaxs = register_scalar_binary( op="pto.tmaxs", name="template_tmaxs", vector_op=pto.vmaxs, - dtypes=same_dtype_signatures(3), + dtypes=_DTYPES, +) + +template_tmaxs_1d = register_scalar_binary( + op="pto.tmaxs", + name="template_tmaxs_1d", + vector_op=pto.vmaxs, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tmin.py b/lib/TileOps/a5/tmin.py index d7fccfcde4..d20bc13c13 100644 --- a/lib/TileOps/a5/tmin.py +++ b/lib/TileOps/a5/tmin.py @@ -8,46 +8,26 @@ """PTODSL TileLib template for pto.tmin (ported from the legacy TileLang template).""" from ptodsl import pto -import ptodsl.tilelib as tilelib -from ._common import NUMERIC_DTYPES +from ._common import same_dtype_signatures +from ._elementwise import register_binary -def _ub_or_vec_row_major(operand_memory_spaces, operand_b_layouts, operand_s_layouts, **_): - return ( - all(space in {"ub", "vec"} for space in operand_memory_spaces) - and all(layout == "row_major" for layout in operand_b_layouts) - and all(layout == "none_box" for layout in operand_s_layouts) - ) +_DTYPES = same_dtype_signatures(3) -@tilelib.tile_template( +template_tmin = register_binary( op="pto.tmin", - target="a5", name="template_tmin", - dtypes=[(dtype, dtype, dtype) for dtype in NUMERIC_DTYPES], - iteration_axis="none", - op_engine="vector", - op_class="elementwise", - constraints=[ - _ub_or_vec_row_major, - tilelib.require_same_valid_shape("src0", "src1", "dst"), - ], - priority=0, - id=0, - loop_depth=2, - is_post_update=False, + vector_op=pto.vmin, + dtypes=_DTYPES, +) + + +template_tmin_1d = register_binary( + op="pto.tmin", + name="template_tmin_1d", + vector_op=pto.vmin, + dtypes=_DTYPES, + traversal="1d", ) -def template_tmin(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - lanes = pto.elements_per_vreg(dtype) - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - min_val = pto.vmin(lhs, rhs, mask) - pto.vsts(min_val, dst[row, col:], mask) diff --git a/lib/TileOps/a5/tmins.py b/lib/TileOps/a5/tmins.py index 214e7998d0..ae73067920 100644 --- a/lib/TileOps/a5/tmins.py +++ b/lib/TileOps/a5/tmins.py @@ -13,9 +13,20 @@ from ._elementwise import register_scalar_binary +_DTYPES = same_dtype_signatures(3) + + template_tmins = register_scalar_binary( op="pto.tmins", name="template_tmins", vector_op=pto.vmins, - dtypes=same_dtype_signatures(3), + dtypes=_DTYPES, +) + +template_tmins_1d = register_scalar_binary( + op="pto.tmins", + name="template_tmins_1d", + vector_op=pto.vmins, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tmul.py b/lib/TileOps/a5/tmul.py index a31e4806d3..e357f460f2 100644 --- a/lib/TileOps/a5/tmul.py +++ b/lib/TileOps/a5/tmul.py @@ -17,9 +17,21 @@ def _vmul(lhs, rhs, mask): return pto.vmul(lhs, rhs, mask) +_DTYPES = same_dtype_signatures(3) + + template_tmul = register_binary( op="pto.tmul", name="template_tmul", vector_op=_vmul, - dtypes=same_dtype_signatures(3), + dtypes=_DTYPES, +) + + +template_tmul_1d = register_binary( + op="pto.tmul", + name="template_tmul_1d", + vector_op=_vmul, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tmuls.py b/lib/TileOps/a5/tmuls.py index f796c0008e..3de9518b49 100644 --- a/lib/TileOps/a5/tmuls.py +++ b/lib/TileOps/a5/tmuls.py @@ -13,9 +13,20 @@ from ._elementwise import register_scalar_binary +_DTYPES = same_dtype_signatures(3) + + template_tmuls = register_scalar_binary( op="pto.tmuls", name="template_tmuls", vector_op=pto.vmuls, - dtypes=same_dtype_signatures(3), + dtypes=_DTYPES, +) + +template_tmuls_1d = register_scalar_binary( + op="pto.tmuls", + name="template_tmuls_1d", + vector_op=pto.vmuls, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tneg.py b/lib/TileOps/a5/tneg.py index f7514dca85..e16eda5682 100644 --- a/lib/TileOps/a5/tneg.py +++ b/lib/TileOps/a5/tneg.py @@ -12,16 +12,28 @@ from ._elementwise import register_unary +_DTYPES = [ + ("i8", "i8"), + ("i16", "i16"), + ("i32", "i32"), + ("f16", "f16"), + ("bf16", "bf16"), + ("f32", "f32"), +] + + template_tneg = register_unary( op="pto.tneg", name="template_tneg", vector_op=pto.vneg, - dtypes=[ - ("i8", "i8"), - ("i16", "i16"), - ("i32", "i32"), - ("f16", "f16"), - ("bf16", "bf16"), - ("f32", "f32"), - ], + dtypes=_DTYPES, +) + + +template_tneg_1d = register_unary( + op="pto.tneg", + name="template_tneg_1d", + vector_op=pto.vneg, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tnot.py b/lib/TileOps/a5/tnot.py index f86792de2b..e21572f0ae 100644 --- a/lib/TileOps/a5/tnot.py +++ b/lib/TileOps/a5/tnot.py @@ -13,9 +13,21 @@ from ._elementwise import register_unary +_DTYPES = [(dtype, dtype) for dtype in INT_DTYPES] + + template_tnot = register_unary( op="pto.tnot", name="template_tnot", vector_op=pto.vnot, - dtypes=[(dtype, dtype) for dtype in INT_DTYPES], + dtypes=_DTYPES, +) + + +template_tnot_1d = register_unary( + op="pto.tnot", + name="template_tnot_1d", + vector_op=pto.vnot, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tor.py b/lib/TileOps/a5/tor.py index c58a4d71c2..18ff1a8a25 100644 --- a/lib/TileOps/a5/tor.py +++ b/lib/TileOps/a5/tor.py @@ -13,9 +13,21 @@ from ._elementwise import register_binary +_DTYPES = [(dtype, dtype, dtype) for dtype in INT_DTYPES] + + template_tor = register_binary( op="pto.tor", name="template_tor", vector_op=pto.vor, - dtypes=[(dtype, dtype, dtype) for dtype in INT_DTYPES], + dtypes=_DTYPES, +) + + +template_tor_1d = register_binary( + op="pto.tor", + name="template_tor_1d", + vector_op=pto.vor, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tors.py b/lib/TileOps/a5/tors.py index b8cece3813..3c28f01fc1 100644 --- a/lib/TileOps/a5/tors.py +++ b/lib/TileOps/a5/tors.py @@ -13,10 +13,22 @@ from ._elementwise import register_scalar_binary +_DTYPES = [(dtype, dtype, dtype) for dtype in INT_DTYPES] + + template_tors = register_scalar_binary( op="pto.tors", name="template_tors", vector_op=pto.vor, broadcast_scalar=True, - dtypes=[(dtype, dtype, dtype) for dtype in INT_DTYPES], + dtypes=_DTYPES, +) + +template_tors_1d = register_scalar_binary( + op="pto.tors", + name="template_tors_1d", + vector_op=pto.vor, + broadcast_scalar=True, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tprelu.py b/lib/TileOps/a5/tprelu.py index 6080beef3d..d9781f6d31 100644 --- a/lib/TileOps/a5/tprelu.py +++ b/lib/TileOps/a5/tprelu.py @@ -11,15 +11,27 @@ from ptodsl import pto +_TPRELU_DTYPES = [ + ("f16", "f16", "f16", "f16"), + ("f32", "f32", "f32", "f32"), + ("f16", "f16", "i8", "f16"), + ("f32", "f32", "i8", "f32"), +] + + template_tprelu = register_binary( op="pto.tprelu", name="template_tprelu", vector_op=pto.vprelu, - dtypes=[ - ("f16", "f16", "f16", "f16"), - ("f32", "f32", "f32", "f32"), - ("f16", "f16", "i8", "f16"), - ("f32", "f32", "i8", "f32"), - ], + dtypes=_TPRELU_DTYPES, + has_tmp=True, +) + +template_tprelu_1d = register_binary( + op="pto.tprelu", + name="template_tprelu_1d", + vector_op=pto.vprelu, + dtypes=_TPRELU_DTYPES, has_tmp=True, + traversal="1d", ) diff --git a/lib/TileOps/a5/trecip.py b/lib/TileOps/a5/trecip.py index 5a28bb586a..409145cd78 100644 --- a/lib/TileOps/a5/trecip.py +++ b/lib/TileOps/a5/trecip.py @@ -10,39 +10,73 @@ from ptodsl import pto import ptodsl.tilelib as tilelib +from ._elementwise import emit_unary_1d, emit_unary_2d, traversal_metadata -@tilelib.tile_template( - op="pto.trecip", - target="a5", - name="template_trecip", - dtypes=[ - ("f16", "f16"), - ("f32", "f32"), - ], - iteration_axis="none", - op_engine="vector", - op_class="elementwise", - constraints=[ + +_DTYPES = [ + ("f16", "f16"), + ("f32", "f32"), +] + + +def _base_constraints(): + return [ tilelib.check_memory_space("ub"), tilelib.check_layout("row_major"), tilelib.check_s_layout("none_box"), - ], - id=0, - loop_depth=2, - is_post_update=False, - tags=("elementwise", "reciprocal"), -) -def template_trecip(src: pto.Tile, dst: pto.Tile): + ] + + +def _emit_trecip(src, dst, traversal): + """Emit the operation-specific reciprocal computation.""" + dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - lanes = pto.elements_per_vreg(dtype) one_scalar = pto.f16(1.0) if str(dtype) == "f16" else pto.f32(1.0) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - value = pto.vlds(src[row, col:]) - one = pto.vbr(one_scalar) - result = pto.vdiv(one, value, mask) - pto.vsts(result, dst[row, col:], mask) + def reciprocal(value, mask): + one = pto.vbr(one_scalar) + return pto.vdiv(one, value, mask) + + if traversal == "1d": + emit_unary_1d(src, dst, reciprocal) + else: + emit_unary_2d(src, dst, reciprocal) + + +def _register_trecip(*, name, traversal): + constraints = _base_constraints() + loop_depth, priority, candidate_id = traversal_metadata(traversal) + if traversal == "1d": + constraints.append(tilelib.require_elementwise_1d("src", "dst")) + + @tilelib.tile_template( + op="pto.trecip", + target="a5", + name=name, + dtypes=_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="elementwise", + constraints=constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, + is_post_update=False, + tags=("elementwise", "reciprocal"), + ) + def template(src: pto.Tile, dst: pto.Tile): + _emit_trecip(src, dst, traversal) + + return template + + +template_trecip = _register_trecip( + name="template_trecip", + traversal="2d", +) + + +template_trecip_1d = _register_trecip( + name="template_trecip_1d", + traversal="1d", +) diff --git a/lib/TileOps/a5/trelu.py b/lib/TileOps/a5/trelu.py index f529eeb880..d630fab298 100644 --- a/lib/TileOps/a5/trelu.py +++ b/lib/TileOps/a5/trelu.py @@ -12,9 +12,21 @@ from ._elementwise import register_unary +_DTYPES = [("i32", "i32"), ("f16", "f16"), ("f32", "f32")] + + template_trelu = register_unary( op="pto.trelu", name="template_trelu", vector_op=pto.vrelu, - dtypes=[("i32", "i32"), ("f16", "f16"), ("f32", "f32")], + dtypes=_DTYPES, +) + + +template_trelu_1d = register_unary( + op="pto.trelu", + name="template_trelu_1d", + vector_op=pto.vrelu, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/trem.py b/lib/TileOps/a5/trem.py index 19d60c83d0..940cbc7d48 100644 --- a/lib/TileOps/a5/trem.py +++ b/lib/TileOps/a5/trem.py @@ -17,3 +17,12 @@ round_mode="F", has_tmp=True, ) + +template_trem_1d = register_binary_remainder( + op="pto.trem", + name="template_trem_1d", + dtypes=REM_DTYPES, + round_mode="F", + has_tmp=True, + traversal="1d", +) diff --git a/lib/TileOps/a5/trems.py b/lib/TileOps/a5/trems.py index 6231af231e..7ce4b56418 100644 --- a/lib/TileOps/a5/trems.py +++ b/lib/TileOps/a5/trems.py @@ -17,3 +17,12 @@ round_mode="F", has_tmp=True, ) + +template_trems_1d = register_scalar_remainder( + op="pto.trems", + name="template_trems_1d", + dtypes=REMS_DTYPES, + round_mode="F", + has_tmp=True, + traversal="1d", +) diff --git a/lib/TileOps/a5/trsqrt.py b/lib/TileOps/a5/trsqrt.py index 9ac78eb3c9..44f711d39e 100644 --- a/lib/TileOps/a5/trsqrt.py +++ b/lib/TileOps/a5/trsqrt.py @@ -12,12 +12,24 @@ from ._elementwise import register_unary +_DTYPES = [ + ("f16", "f16"), + ("f32", "f32"), +] + + template_trsqrt = register_unary( op="pto.trsqrt", name="template_trsqrt", vector_op=pto.vrsqrt, - dtypes=[ - ("f16", "f16"), - ("f32", "f32"), - ], + dtypes=_DTYPES, +) + + +template_trsqrt_1d = register_unary( + op="pto.trsqrt", + name="template_trsqrt_1d", + vector_op=pto.vrsqrt, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tsel.py b/lib/TileOps/a5/tsel.py index 017835401d..564ff26870 100644 --- a/lib/TileOps/a5/tsel.py +++ b/lib/TileOps/a5/tsel.py @@ -5,36 +5,105 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""PTODSL TileLib template for pto.tsel.""" +"""PTODSL TileLib templates for pto.tsel.""" from ptodsl import pto +from ptodsl._ast_rewrite import rewrite_jit_function import ptodsl.tilelib as tilelib from ._common import ub_row_major_constraints +from ._elementwise import traversal_metadata -@tilelib.tile_template( - op="pto.tsel", - target="a5", - name="template_tsel", - dtypes=[ - ("i8", "f32", "f32", "f32", "f32"), - ("i8", "f16", "f16", "f16", "f16"), - ("i8", "i8", "i8", "i8", "i8"), - ], - iteration_axis="none", - op_engine="vector", - op_class="other", - constraints=ub_row_major_constraints( - "src0", "src1", "tmp", "dst", require_same_valid_shape=False - ), - id=0, - loop_depth=2, - is_post_update=False, - tags=("select", "predicate-load"), -) -def template_tsel(mask: pto.Tile, src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - _ = tmp +_DTYPES = [ + ("i8", "f32", "f32", "f32", "f32"), + ("i8", "f16", "f16", "f16", "f16"), + ("i8", "i8", "i8", "i8", "i8"), +] + + +def _f32_select_masks(raw_mask, full_mask_b16): + select_mask = pto.pbitcast(raw_mask, pto.mask_b16) + select_mask0, select_mask1 = pto.pintlv_b16( + select_mask, + full_mask_b16, + ) + return ( + pto.pbitcast(select_mask0, pto.mask_b32), + pto.pbitcast(select_mask1, pto.mask_b32), + ) + + +def _f32_paired_cols(valid_cols, lanes): + repeat_times = (valid_cols + lanes - 1) // lanes + return (repeat_times // 2) * lanes * 2 + + +@rewrite_jit_function +def _emit_tsel_1d(mask, src0, src1, dst): + dtype = dst.dtype + valid_rows, valid_cols = dst.valid_shape + lanes = pto.elements_per_vreg(dtype) + total_elements = valid_rows * valid_cols + mask_ptr = pto.castptr(mask.as_ptr(), pto.ptr(pto.ui8, "ub")) + src0_ptr = src0.as_ptr() + src1_ptr = src1.as_ptr() + dst_ptr = dst.as_ptr() + + if str(dtype) == "f32": + full_mask_b16 = pto.pset_b16(pto.PAT.ALL) + remained = total_elements + for offset in range(0, total_elements, lanes * 2): + raw_mask = pto.plds( + mask_ptr, + offset // 8, + dist=pto.PredicateDist.US, + ) + select_mask0, select_mask1 = _f32_select_masks( + raw_mask, + full_mask_b16, + ) + pred0, remained = pto.make_mask(dtype, remained) + pred1, remained = pto.make_mask(dtype, remained) + lhs0 = pto.vlds(src0_ptr, offset) + rhs0 = pto.vlds(src1_ptr, offset) + lhs1 = pto.vlds(src0_ptr, offset + lanes) + rhs1 = pto.vlds(src1_ptr, offset + lanes) + selected0 = pto.vsel(lhs0, rhs0, select_mask0) + selected1 = pto.vsel(lhs1, rhs1, select_mask1) + pto.vsts(selected0, dst_ptr, offset, pred0) + pto.vsts(selected1, dst_ptr, offset + lanes, pred1) + elif str(dtype) == "f16": + remained = total_elements + for offset in range(0, total_elements, lanes): + pred, remained = pto.make_mask(dtype, remained) + select_mask = pto.plds( + mask_ptr, + offset // 8, + dist=pto.PredicateDist.US, + ) + select_mask = pto.pbitcast(select_mask, pto.mask_b16) + lhs = pto.vlds(src0_ptr, offset) + rhs = pto.vlds(src1_ptr, offset) + result = pto.vsel(lhs, rhs, select_mask) + pto.vsts(result, dst_ptr, offset, pred) + else: + remained = total_elements + for offset in range(0, total_elements, lanes): + pred, remained = pto.make_mask(dtype, remained) + select_mask = pto.plds( + mask_ptr, + offset // 8, + dist=pto.PredicateDist.NORM, + ) + lhs = pto.vlds(src0_ptr, offset) + rhs = pto.vlds(src1_ptr, offset) + result = pto.vsel(lhs, rhs, select_mask) + pto.vsts(result, dst_ptr, offset, pred) + + +@rewrite_jit_function +def _emit_tsel_2d(mask, src0, src1, dst): dtype = dst.dtype valid_rows, valid_cols = dst.valid_shape lanes = pto.elements_per_vreg(dtype) @@ -44,17 +113,22 @@ def template_tsel(mask: pto.Tile, src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, if str(dtype) == "f32": full_mask_b16 = pto.pset_b16(pto.PAT.ALL) pair_width = lanes * 2 - paired_cols = (valid_cols // pair_width) * pair_width + paired_cols = _f32_paired_cols(valid_cols, lanes) for row in range(0, valid_rows, 1): + remained = valid_cols for col in range(0, paired_cols, pair_width): mask_offset = row * mask_stride + col // 8 - select_mask_raw = pto.plds(mask_ptr, mask_offset, dist=pto.PredicateDist.US) - select_mask = pto.pbitcast(select_mask_raw, pto.mask_b16) - pred0, _ = pto.make_mask(dtype, pair_width) - pred1, _ = pto.make_mask(dtype, lanes) - select_mask0, select_mask1 = pto.pintlv_b16(select_mask, full_mask_b16) - select_mask0 = pto.pbitcast(select_mask0, pto.mask_b32) - select_mask1 = pto.pbitcast(select_mask1, pto.mask_b32) + raw_mask = pto.plds( + mask_ptr, + mask_offset, + dist=pto.PredicateDist.US, + ) + select_mask0, select_mask1 = _f32_select_masks( + raw_mask, + full_mask_b16, + ) + pred0, remained = pto.make_mask(dtype, remained) + pred1, remained = pto.make_mask(dtype, remained) lhs0 = pto.vlds(src0[row, col:]) rhs0 = pto.vlds(src1[row, col:]) lhs1 = pto.vlds(src0[row, col + lanes:]) @@ -63,26 +137,36 @@ def template_tsel(mask: pto.Tile, src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, selected1 = pto.vsel(lhs1, rhs1, select_mask1) pto.vsts(selected0, dst[row, col:], pred0) pto.vsts(selected1, dst[row, col + lanes:], pred1) - tail_cols = valid_cols - paired_cols - if tail_cols > 0: - col = paired_cols + + for col in range(paired_cols, valid_cols, lanes): mask_offset = row * mask_stride + col // 8 - select_mask_raw = pto.plds(mask_ptr, mask_offset, dist=pto.PredicateDist.US) - select_mask = pto.pbitcast(select_mask_raw, pto.mask_b16) - select_mask0 = pto.punpack(select_mask, pto.PredicatePart.LOWER) - select_mask0 = pto.pbitcast(select_mask0, pto.mask_b32) - pred0, _ = pto.make_mask(dtype, tail_cols) - lhs0 = pto.vlds(src0[row, col:]) - rhs0 = pto.vlds(src1[row, col:]) - selected0 = pto.vsel(lhs0, rhs0, select_mask0) - pto.vsts(selected0, dst[row, col:], pred0) + raw_mask = pto.plds( + mask_ptr, + mask_offset, + dist=pto.PredicateDist.US, + ) + select_mask = pto.pbitcast(raw_mask, pto.mask_b16) + select_mask = pto.punpack( + select_mask, + pto.PredicatePart.LOWER, + ) + select_mask = pto.pbitcast(select_mask, pto.mask_b32) + pred, remained = pto.make_mask(dtype, remained) + lhs = pto.vlds(src0[row, col:]) + rhs = pto.vlds(src1[row, col:]) + selected = pto.vsel(lhs, rhs, select_mask) + pto.vsts(selected, dst[row, col:], pred) elif str(dtype) == "f16": for row in range(0, valid_rows, 1): remained = valid_cols for col in range(0, valid_cols, lanes): pred, remained = pto.make_mask(dtype, remained) mask_offset = row * mask_stride + col // 8 - select_mask = pto.plds(mask_ptr, mask_offset, dist=pto.PredicateDist.US) + select_mask = pto.plds( + mask_ptr, + mask_offset, + dist=pto.PredicateDist.US, + ) select_mask = pto.pbitcast(select_mask, pto.mask_b16) lhs = pto.vlds(src0[row, col:]) rhs = pto.vlds(src1[row, col:]) @@ -94,8 +178,75 @@ def template_tsel(mask: pto.Tile, src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, for col in range(0, valid_cols, lanes): pred, remained = pto.make_mask(dtype, remained) mask_offset = row * mask_stride + col // 8 - select_mask = pto.plds(mask_ptr, mask_offset, dist=pto.PredicateDist.NORM) + select_mask = pto.plds( + mask_ptr, + mask_offset, + dist=pto.PredicateDist.NORM, + ) lhs = pto.vlds(src0[row, col:]) rhs = pto.vlds(src1[row, col:]) result = pto.vsel(lhs, rhs, select_mask) pto.vsts(result, dst[row, col:], pred) + + +def _register_tsel(*, name, traversal): + constraints = ub_row_major_constraints( + "src0", + "src1", + "tmp", + "dst", + require_same_valid_shape=False, + ) + loop_depth, priority, candidate_id = traversal_metadata(traversal) + if traversal == "1d": + constraints.append( + tilelib.require_predicate_select_1d( + "mask", + "src0", + "src1", + "dst", + temporary_operand="tmp", + memory_spaces=("ub",), + ) + ) + + @tilelib.tile_template( + op="pto.tsel", + target="a5", + name=name, + dtypes=_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="other", + constraints=constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, + is_post_update=False, + tags=("select", "predicate-load"), + ) + def template( + mask: pto.Tile, + src0: pto.Tile, + src1: pto.Tile, + tmp: pto.Tile, + dst: pto.Tile, + ): + _ = tmp + if traversal == "1d": + _emit_tsel_1d(mask, src0, src1, dst) + else: + _emit_tsel_2d(mask, src0, src1, dst) + + return template + + +template_tsel = _register_tsel( + name="template_tsel", + traversal="2d", +) + +template_tsel_1d = _register_tsel( + name="template_tsel_1d", + traversal="1d", +) diff --git a/lib/TileOps/a5/tsels.py b/lib/TileOps/a5/tsels.py index 62ccef8e86..9268961934 100644 --- a/lib/TileOps/a5/tsels.py +++ b/lib/TileOps/a5/tsels.py @@ -5,104 +5,196 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""PTODSL TileLib template for pto.tsels.""" +"""PTODSL TileLib templates for pto.tsels.""" from ptodsl import pto +from ptodsl._ast_rewrite import rewrite_jit_function import ptodsl.tilelib as tilelib +from ._elementwise import traversal_metadata -def _tsels_shapes(mask_valid_shape=(), src_valid_shape=(), tmp_valid_shape=(), dst_valid_shape=(), **_): + +_DTYPES = [ + ("i8", "i8", "i8", "i8", "i8"), + ("i16", "i8", "i8", "i8", "i8"), + ("i32", "i8", "i8", "i8", "i8"), + ("i8", "i16", "i16", "i16", "i16"), + ("i16", "i16", "i16", "i16", "i16"), + ("i32", "i16", "i16", "i16", "i16"), + ("i8", "i32", "i32", "i32", "i32"), + ("i16", "i32", "i32", "i32", "i32"), + ("i32", "i32", "i32", "i32", "i32"), + ("i8", "f32", "f32", "f32", "f32"), + ("i16", "f32", "f32", "f32", "f32"), + ("i32", "f32", "f32", "f32", "f32"), + ("i8", "f16", "f16", "f16", "f16"), + ("i16", "f16", "f16", "f16", "f16"), + ("i32", "f16", "f16", "f16", "f16"), +] + + +def _tsels_shapes( + mask_valid_shape=(), + src_valid_shape=(), + tmp_valid_shape=(), + dst_valid_shape=(), + **_, +): _ = mask_valid_shape, tmp_valid_shape - return len(src_valid_shape) == 2 and tuple(src_valid_shape) == tuple(dst_valid_shape) + return ( + len(src_valid_shape) == 2 + and tuple(src_valid_shape) == tuple(dst_valid_shape) + ) -@tilelib.tile_template( - op="pto.tsels", - target="a5", - name="template_tsels", - dtypes=[ - ("i8", "i8", "i8", "i8", "i8"), - ("i16", "i8", "i8", "i8", "i8"), - ("i32", "i8", "i8", "i8", "i8"), - ("i8", "i16", "i16", "i16", "i16"), - ("i16", "i16", "i16", "i16", "i16"), - ("i32", "i16", "i16", "i16", "i16"), - ("i8", "i32", "i32", "i32", "i32"), - ("i16", "i32", "i32", "i32", "i32"), - ("i32", "i32", "i32", "i32", "i32"), - ("i8", "f32", "f32", "f32", "f32"), - ("i16", "f32", "f32", "f32", "f32"), - ("i32", "f32", "f32", "f32", "f32"), - ("i8", "f16", "f16", "f16", "f16"), - ("i16", "f16", "f16", "f16", "f16"), - ("i32", "f16", "f16", "f16", "f16"), - ], - iteration_axis="none", - op_engine="vector", - op_class="other", - constraints=[ - tilelib.check_memory_space("ub"), - tilelib.check_layout("row_major"), - tilelib.check_s_layout("none_box"), - _tsels_shapes, - ], - id=0, - loop_depth=2, - is_post_update=False, - tags=("select", "scalar", "predicate-load"), -) -def template_tsels(mask: pto.Tile, src: pto.Tile, tmp: pto.Tile, scalar, dst: pto.Tile): - _ = tmp +def _f32_select_masks(raw_mask, full_mask_b16): + select_mask = pto.pbitcast(raw_mask, pto.mask_b16) + select_mask0, select_mask1 = pto.pintlv_b16( + select_mask, + full_mask_b16, + ) + return ( + pto.pbitcast(select_mask0, pto.mask_b32), + pto.pbitcast(select_mask1, pto.mask_b32), + ) + + +def _f32_paired_cols(valid_cols, lanes): + repeat_times = (valid_cols + lanes - 1) // lanes + return (repeat_times // 2) * lanes * 2 + + +def _scalar_vector(dtype, scalar): + lanes = pto.elements_per_vreg(dtype) + full_mask, _ = pto.make_mask(dtype, lanes) + return pto.vdup(scalar, full_mask) + + +@rewrite_jit_function +def _emit_tsels_1d(mask, src, scalar, dst): + dtype = dst.dtype + valid_rows, valid_cols = dst.valid_shape + lanes = pto.elements_per_vreg(dtype) + total_elements = valid_rows * valid_cols + mask_ptr = pto.castptr(mask.as_ptr(), pto.ptr(pto.ui8, "ub")) + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + scalar_vec = _scalar_vector(dtype, scalar) + + if str(dtype) in {"f32", "i32"}: + full_mask_b16 = pto.pset_b16(pto.PAT.ALL) + remained = total_elements + for offset in range(0, total_elements, lanes * 2): + raw_mask = pto.plds( + mask_ptr, + offset // 8, + dist=pto.PredicateDist.US, + ) + select_mask0, select_mask1 = _f32_select_masks( + raw_mask, + full_mask_b16, + ) + pred0, remained = pto.make_mask(dtype, remained) + pred1, remained = pto.make_mask(dtype, remained) + value0 = pto.vlds(src_ptr, offset) + value1 = pto.vlds(src_ptr, offset + lanes) + selected0 = pto.vsel(value0, scalar_vec, select_mask0) + selected1 = pto.vsel(value1, scalar_vec, select_mask1) + pto.vsts(selected0, dst_ptr, offset, pred0) + pto.vsts(selected1, dst_ptr, offset + lanes, pred1) + elif str(dtype) in {"f16", "i16"}: + remained = total_elements + for offset in range(0, total_elements, lanes): + pred, remained = pto.make_mask(dtype, remained) + select_mask = pto.plds( + mask_ptr, + offset // 8, + dist=pto.PredicateDist.US, + ) + select_mask = pto.pbitcast(select_mask, pto.mask_b16) + value = pto.vlds(src_ptr, offset) + result = pto.vsel(value, scalar_vec, select_mask) + pto.vsts(result, dst_ptr, offset, pred) + else: + remained = total_elements + for offset in range(0, total_elements, lanes): + pred, remained = pto.make_mask(dtype, remained) + select_mask = pto.plds( + mask_ptr, + offset // 8, + dist=pto.PredicateDist.NORM, + ) + value = pto.vlds(src_ptr, offset) + result = pto.vsel(value, scalar_vec, select_mask) + pto.vsts(result, dst_ptr, offset, pred) + + +@rewrite_jit_function +def _emit_tsels_2d(mask, src, scalar, dst): dtype = dst.dtype valid_rows, valid_cols = dst.valid_shape lanes = pto.elements_per_vreg(dtype) mask_ptr = pto.castptr(mask.as_ptr(), pto.ptr(pto.ui8, "ub")) mask_stride = mask.shape[1] * pto.bytewidth(mask.dtype) - full_mask, _ = pto.make_mask(dtype, lanes) - scalar_vec = pto.vdup(scalar, full_mask) + scalar_vec = _scalar_vector(dtype, scalar) - if lanes == 64: + if str(dtype) in {"f32", "i32"}: full_mask_b16 = pto.pset_b16(pto.PAT.ALL) pair_width = lanes * 2 - paired_cols = (valid_cols // pair_width) * pair_width + paired_cols = _f32_paired_cols(valid_cols, lanes) for row in range(0, valid_rows, 1): + remained = valid_cols for col in range(0, paired_cols, pair_width): mask_offset = row * mask_stride + col // 8 - select_mask_raw = pto.plds(mask_ptr, mask_offset, dist=pto.PredicateDist.US) - select_mask = pto.pbitcast(select_mask_raw, pto.mask_b16) - pred0, _ = pto.make_mask(dtype, pair_width) - pred1, _ = pto.make_mask(dtype, lanes) - select_mask0, select_mask1 = pto.pintlv_b16(select_mask, full_mask_b16) - select_mask0 = pto.pbitcast(select_mask0, pto.mask_b32) - select_mask1 = pto.pbitcast(select_mask1, pto.mask_b32) - src0 = pto.vlds(src[row, col:]) - src1 = pto.vlds(src[row, col + lanes:]) - selected0 = pto.vsel(src0, scalar_vec, select_mask0) - selected1 = pto.vsel(src1, scalar_vec, select_mask1) + raw_mask = pto.plds( + mask_ptr, + mask_offset, + dist=pto.PredicateDist.US, + ) + select_mask0, select_mask1 = _f32_select_masks( + raw_mask, + full_mask_b16, + ) + pred0, remained = pto.make_mask(dtype, remained) + pred1, remained = pto.make_mask(dtype, remained) + value0 = pto.vlds(src[row, col:]) + value1 = pto.vlds(src[row, col + lanes:]) + selected0 = pto.vsel(value0, scalar_vec, select_mask0) + selected1 = pto.vsel(value1, scalar_vec, select_mask1) pto.vsts(selected0, dst[row, col:], pred0) pto.vsts(selected1, dst[row, col + lanes:], pred1) - tail_cols = valid_cols - paired_cols - if tail_cols > 0: - col = paired_cols + + for col in range(paired_cols, valid_cols, lanes): mask_offset = row * mask_stride + col // 8 - select_mask_raw = pto.plds(mask_ptr, mask_offset, dist=pto.PredicateDist.US) - select_mask = pto.pbitcast(select_mask_raw, pto.mask_b16) - select_mask0 = pto.punpack(select_mask, pto.PredicatePart.LOWER) - select_mask0 = pto.pbitcast(select_mask0, pto.mask_b32) - pred0, _ = pto.make_mask(dtype, tail_cols) - src0 = pto.vlds(src[row, col:]) - selected0 = pto.vsel(src0, scalar_vec, select_mask0) - pto.vsts(selected0, dst[row, col:], pred0) - elif lanes == 128: + raw_mask = pto.plds( + mask_ptr, + mask_offset, + dist=pto.PredicateDist.US, + ) + select_mask = pto.pbitcast(raw_mask, pto.mask_b16) + select_mask = pto.punpack( + select_mask, + pto.PredicatePart.LOWER, + ) + select_mask = pto.pbitcast(select_mask, pto.mask_b32) + pred, remained = pto.make_mask(dtype, remained) + value = pto.vlds(src[row, col:]) + selected = pto.vsel(value, scalar_vec, select_mask) + pto.vsts(selected, dst[row, col:], pred) + elif str(dtype) in {"f16", "i16"}: for row in range(0, valid_rows, 1): remained = valid_cols for col in range(0, valid_cols, lanes): pred, remained = pto.make_mask(dtype, remained) mask_offset = row * mask_stride + col // 8 - select_mask = pto.plds(mask_ptr, mask_offset, dist=pto.PredicateDist.US) + select_mask = pto.plds( + mask_ptr, + mask_offset, + dist=pto.PredicateDist.US, + ) select_mask = pto.pbitcast(select_mask, pto.mask_b16) - lhs = pto.vlds(src[row, col:]) - result = pto.vsel(lhs, scalar_vec, select_mask) + value = pto.vlds(src[row, col:]) + result = pto.vsel(value, scalar_vec, select_mask) pto.vsts(result, dst[row, col:], pred) else: for row in range(0, valid_rows, 1): @@ -110,7 +202,72 @@ def template_tsels(mask: pto.Tile, src: pto.Tile, tmp: pto.Tile, scalar, dst: pt for col in range(0, valid_cols, lanes): pred, remained = pto.make_mask(dtype, remained) mask_offset = row * mask_stride + col // 8 - select_mask = pto.plds(mask_ptr, mask_offset, dist=pto.PredicateDist.NORM) - lhs = pto.vlds(src[row, col:]) - result = pto.vsel(lhs, scalar_vec, select_mask) + select_mask = pto.plds( + mask_ptr, + mask_offset, + dist=pto.PredicateDist.NORM, + ) + value = pto.vlds(src[row, col:]) + result = pto.vsel(value, scalar_vec, select_mask) pto.vsts(result, dst[row, col:], pred) + + +def _register_tsels(*, name, traversal): + constraints = [ + tilelib.check_memory_space("ub"), + tilelib.check_layout("row_major"), + tilelib.check_s_layout("none_box"), + _tsels_shapes, + ] + loop_depth, priority, candidate_id = traversal_metadata(traversal) + if traversal == "1d": + constraints.append( + tilelib.require_predicate_select_1d( + "mask", + "src", + "dst", + temporary_operand="tmp", + memory_spaces=("ub",), + ) + ) + + @tilelib.tile_template( + op="pto.tsels", + target="a5", + name=name, + dtypes=_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="other", + constraints=constraints, + priority=priority, + id=candidate_id, + loop_depth=loop_depth, + is_post_update=False, + tags=("select", "scalar", "predicate-load"), + ) + def template( + mask: pto.Tile, + src: pto.Tile, + tmp: pto.Tile, + scalar, + dst: pto.Tile, + ): + _ = tmp + if traversal == "1d": + _emit_tsels_1d(mask, src, scalar, dst) + else: + _emit_tsels_2d(mask, src, scalar, dst) + + return template + + +template_tsels = _register_tsels( + name="template_tsels", + traversal="2d", +) + +template_tsels_1d = _register_tsels( + name="template_tsels_1d", + traversal="1d", +) diff --git a/lib/TileOps/a5/tshl.py b/lib/TileOps/a5/tshl.py index 8d4a5ce5ba..3342313beb 100644 --- a/lib/TileOps/a5/tshl.py +++ b/lib/TileOps/a5/tshl.py @@ -13,9 +13,21 @@ from ._elementwise import register_binary +_DTYPES = [(dtype, dtype, dtype) for dtype in INT_DTYPES] + + template_tshl = register_binary( op="pto.tshl", name="template_tshl", vector_op=pto.vshl, - dtypes=[(dtype, dtype, dtype) for dtype in INT_DTYPES], + dtypes=_DTYPES, +) + + +template_tshl_1d = register_binary( + op="pto.tshl", + name="template_tshl_1d", + vector_op=pto.vshl, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tshls.py b/lib/TileOps/a5/tshls.py index 1f4804b2af..e1d80b6f68 100644 --- a/lib/TileOps/a5/tshls.py +++ b/lib/TileOps/a5/tshls.py @@ -13,9 +13,20 @@ from ._elementwise import register_scalar_binary +_DTYPES = [(dtype, "i16", dtype) for dtype in INT_DTYPES] + + template_tshls = register_scalar_binary( op="pto.tshls", name="template_tshls", vector_op=pto.vshls, - dtypes=[(dtype, "i16", dtype) for dtype in INT_DTYPES], + dtypes=_DTYPES, +) + +template_tshls_1d = register_scalar_binary( + op="pto.tshls", + name="template_tshls_1d", + vector_op=pto.vshls, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tshr.py b/lib/TileOps/a5/tshr.py index 91de7f64ef..000c20e810 100644 --- a/lib/TileOps/a5/tshr.py +++ b/lib/TileOps/a5/tshr.py @@ -13,9 +13,21 @@ from ._elementwise import register_binary +_DTYPES = [(dtype, dtype, dtype) for dtype in INT_DTYPES] + + template_tshr = register_binary( op="pto.tshr", name="template_tshr", vector_op=pto.vshr, - dtypes=[(dtype, dtype, dtype) for dtype in INT_DTYPES], + dtypes=_DTYPES, +) + + +template_tshr_1d = register_binary( + op="pto.tshr", + name="template_tshr_1d", + vector_op=pto.vshr, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tshrs.py b/lib/TileOps/a5/tshrs.py index 900ebcbe5b..a1862ecb56 100644 --- a/lib/TileOps/a5/tshrs.py +++ b/lib/TileOps/a5/tshrs.py @@ -13,9 +13,20 @@ from ._elementwise import register_scalar_binary +_DTYPES = [(dtype, "i16", dtype) for dtype in INT_DTYPES] + + template_tshrs = register_scalar_binary( op="pto.tshrs", name="template_tshrs", vector_op=pto.vshrs, - dtypes=[(dtype, "i16", dtype) for dtype in INT_DTYPES], + dtypes=_DTYPES, +) + +template_tshrs_1d = register_scalar_binary( + op="pto.tshrs", + name="template_tshrs_1d", + vector_op=pto.vshrs, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tsqrt.py b/lib/TileOps/a5/tsqrt.py index 81d2653e5c..99f734abc3 100644 --- a/lib/TileOps/a5/tsqrt.py +++ b/lib/TileOps/a5/tsqrt.py @@ -12,12 +12,24 @@ from ._elementwise import register_unary +_DTYPES = [ + ("f16", "f16"), + ("f32", "f32"), +] + + template_tsqrt = register_unary( op="pto.tsqrt", name="template_tsqrt", vector_op=pto.vsqrt, - dtypes=[ - ("f16", "f16"), - ("f32", "f32"), - ], + dtypes=_DTYPES, +) + + +template_tsqrt_1d = register_unary( + op="pto.tsqrt", + name="template_tsqrt_1d", + vector_op=pto.vsqrt, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tsub.py b/lib/TileOps/a5/tsub.py index 36c34e4c08..2a29ecff34 100644 --- a/lib/TileOps/a5/tsub.py +++ b/lib/TileOps/a5/tsub.py @@ -17,9 +17,21 @@ def _vsub(lhs, rhs, mask): return pto.vsub(lhs, rhs, mask) +_DTYPES = same_dtype_signatures(3) + + template_tsub = register_binary( op="pto.tsub", name="template_tsub", vector_op=_vsub, - dtypes=same_dtype_signatures(3), + dtypes=_DTYPES, +) + + +template_tsub_1d = register_binary( + op="pto.tsub", + name="template_tsub_1d", + vector_op=_vsub, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/lib/TileOps/a5/tsubs.py b/lib/TileOps/a5/tsubs.py index 1bf3dfe2f3..fb7fb48592 100644 --- a/lib/TileOps/a5/tsubs.py +++ b/lib/TileOps/a5/tsubs.py @@ -8,36 +8,27 @@ """PTODSL TileLib template for pto.tsubs.""" from ptodsl import pto -import ptodsl.tilelib as tilelib from ._common import same_dtype_signatures -from ._elementwise import _common_constraints +from ._elementwise import register_scalar_binary -@tilelib.tile_template( +_DTYPES = same_dtype_signatures(3) + + +template_tsubs = register_scalar_binary( op="pto.tsubs", - target="a5", name="template_tsubs", - dtypes=same_dtype_signatures(3), - iteration_axis="none", - op_engine="vector", - op_class="elementwise", - constraints=_common_constraints("src", "dst"), - id=0, - loop_depth=2, - is_post_update=False, - tags=("elementwise", "scalar"), + vector_op=pto.vsub, + broadcast_scalar=True, + dtypes=_DTYPES, ) -def template_tsubs(src: pto.Tile, scalar, dst: pto.Tile): - dtype = dst.dtype - valid_rows, valid_cols = dst.valid_shape - lanes = pto.elements_per_vreg(dtype) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - value = pto.vlds(src[row, col:]) - scalar_value = pto.vbr(scalar) - result = pto.vsub(value, scalar_value, mask) - pto.vsts(result, dst[row, col:], mask) +template_tsubs_1d = register_scalar_binary( + op="pto.tsubs", + name="template_tsubs_1d", + vector_op=pto.vsub, + broadcast_scalar=True, + dtypes=_DTYPES, + traversal="1d", +) diff --git a/lib/TileOps/a5/txor.py b/lib/TileOps/a5/txor.py index 0a134576c7..305cde0c64 100644 --- a/lib/TileOps/a5/txor.py +++ b/lib/TileOps/a5/txor.py @@ -13,10 +13,25 @@ from ._elementwise import register_binary +_TXOR_DTYPES = [ + (dtype, dtype, dtype, dtype) + for dtype in INT_DTYPES +] + + template_txor = register_binary( op="pto.txor", name="template_txor", vector_op=pto.vxor, - dtypes=[(dtype, dtype, dtype, dtype) for dtype in INT_DTYPES], + dtypes=_TXOR_DTYPES, + has_tmp=True, +) + +template_txor_1d = register_binary( + op="pto.txor", + name="template_txor_1d", + vector_op=pto.vxor, + dtypes=_TXOR_DTYPES, has_tmp=True, + traversal="1d", ) diff --git a/lib/TileOps/a5/txors.py b/lib/TileOps/a5/txors.py index 89f298ef64..413194e213 100644 --- a/lib/TileOps/a5/txors.py +++ b/lib/TileOps/a5/txors.py @@ -13,6 +13,12 @@ from ._elementwise import register_scalar_binary +_DTYPES = [ + (dtype, dtype, dtype, dtype) + for dtype in INT_DTYPES +] + + template_txors = register_scalar_binary( op="pto.txors", name="template_txors", @@ -20,5 +26,16 @@ broadcast_scalar=True, has_tmp=True, tmp_matches_src_dst=False, - dtypes=[(dtype, dtype, dtype, dtype) for dtype in INT_DTYPES], + dtypes=_DTYPES, +) + +template_txors_1d = register_scalar_binary( + op="pto.txors", + name="template_txors_1d", + vector_op=pto.vxor, + broadcast_scalar=True, + has_tmp=True, + tmp_matches_src_dst=False, + dtypes=_DTYPES, + traversal="1d", ) diff --git a/ptodsl/docs/developer_guide/tilelib-template-authoring.md b/ptodsl/docs/developer_guide/tilelib-template-authoring.md index a2a418dafa..ac52a53216 100644 --- a/ptodsl/docs/developer_guide/tilelib-template-authoring.md +++ b/ptodsl/docs/developer_guide/tilelib-template-authoring.md @@ -66,8 +66,8 @@ errors are faster to read than opaque constraint failures. Template parameters receive concrete specs built from MLIR operands: -- `TileSpec`: rank-2 tile shape, dtype, memory space, valid shape, layouts, and - pad value. +- `TileSpec`: rank-2 tile shape, dtype, memory space, valid shape, layouts, + fractal size, pad value, and compact mode. - `ViewSpec`: view shape, dtype, memory space, optional strides, and optional layout. - `ScalarSpec`: scalar dtype plus static integer value when known. @@ -95,6 +95,41 @@ Keep predicates small and named by the rule they enforce. A good predicate answers one question, such as "is this a row-major vec tile" or "does this view have a static stride". +Ordinary element-wise 1D candidates should use +`require_elementwise_1d(*operand_names)` and explicitly name every traversed +tile, including temporary operands. It proves a common static logical range on +gap-free row-major/none-box local storage. Predicate-mask and dtype-width- +changing conversion families need additional representation-specific rules. + +Packed compare candidates use +`require_predicate_compare_1d(*data_operand_names, +predicate_operand="dst")`. The data operands still need a common contiguous +logical range, while the predicate destination is checked as one bit per +source element with dtype-dependent complete-load/store rounding. Multi-row +flattening additionally requires every source row to end on a complete +predicate-store boundary and the predicate destination row stride to contain +exactly that row's packed bytes; its physical row must retain the target's +32-byte alignment. Do not substitute ordinary same-shape legality for this +rule: padded predicate rows must restart in the 2D fallback. + +Packed select candidates use +`require_predicate_select_1d(predicate_operand, *data_operand_names, +temporary_operand=...)`. Data operands follow ordinary continuity and +same-range rules, while the mask is checked in physical bytes because `tsels` +allows i8, i16, and i32 mask storage containers for the same packed bits. +Multi-row masks must have no predicate-row padding. On A5 the select temporary +is unused, so validate that its local tile metadata is complete and supported +without requiring its shape to match the data range. + +Width-changing conversion candidates use `require_conversion_1d()`. The rule +proves that source and destination are independently gap-free typed streams; +ordinary conversions have equal element shapes, while packed forms pass an +explicit source-elements-per-destination ratio. Keep unpack/pack distribution +modes, mask dtypes, rounding, saturation, and multi-step instruction sequences +inside the conversion module. Unknown metadata falls back to 2D. If a future +conversion has a tile temporary, extend the family rule to validate it rather +than applying the current source/destination-only predicate unchanged. + When porting a legacy TileLangDSL version, first write down which rule made the legacy version legal. Then encode that rule as metadata or a predicate. Avoid fixing a single failing case by weakening a predicate beyond the legacy @@ -128,6 +163,8 @@ For multi-candidate ops: - assign stable `id` values; - keep ids unique for the op; +- use `priority` rather than ID to express preference; +- avoid equal top-priority candidates for the same concrete operands; - keep names descriptive enough for IR dumps; - add a lit check when candidate ordering matters. @@ -153,18 +190,156 @@ rule that matches TileLangDSL behavior. Template bodies execute under PTODSL tracing. Python values and PTODSL runtime values are not interchangeable. -Use PTODSL control-flow and scalar APIs when a value is runtime-dependent. +Source-backed TileLib template functions use PTODSL's control-flow AST rewrite. +Plain Python `if` and `for ... in range(...)` in a template body or its nested +helpers therefore lower to runtime structured control flow. Prefer that syntax +for ordinary loops: + +```python +remained = valid_cols +for col in range(0, valid_cols, lanes): + mask, remained = pto.make_mask(dtype, remained) + ... +``` + +Use `pto.static_range(...)` when a loop should execute during tracing. Keep the +explicit `pto.if_` and `pto.for_` APIs for unsupported or deliberately explicit +control-flow patterns. + +Module-level Python helpers are outside the registered template function's +source tree, so their bodies are not rewritten merely because the template +calls them. A module-level helper that contains runtime Python control flow +must opt into `rewrite_jit_function`, as the shared element-wise traversal +cores do, or use the explicit control-flow APIs. + Avoid: -- native Python `if` on a PTODSL runtime value; -- native Python `range` using runtime bounds; - assigning Python integers into runtime branch state; - assuming a scalar operand has a compile-time value unless `ScalarSpec.value` - is present. + is present; +- relying on native runtime control-flow rewrite for functions without + retrievable source. This class of bug often appears after selection is fixed: the template becomes legal, then tracing fails in a larger non-smoke path. +## Shared Element-wise Traversal Forms + +Ordinary A5 unary, Tile-Tile, Tile-Scalar, and scalar-fill candidates should +use the registration helpers in `templates/a5/_elementwise.py`. Each registrar +accepts `traversal="1d"` or `traversal="2d"`: + +- `1d` derives `loop_depth=1`, adds + `require_elementwise_1d(*operand_names)`, and emits one vector loop over + `valid_rows * valid_cols`; +- `2d` derives `loop_depth=2` and emits the general row loop plus per-row + vector loop. + +The default remains `2d`, so adding the shared foundation does not silently +change existing operation selection. A migrated operation normally preserves +its existing candidate as the 2D fallback and adds a separately named, +higher-priority 1D candidate: + +```python +template_tadd = register_binary( + op="pto.tadd", + name="template_tadd", + vector_op=pto.vadd, + dtypes=DTYPES, + traversal="2d", +) + +template_tadd_1d = register_binary( + op="pto.tadd", + name="template_tadd_1d", + vector_op=pto.vadd, + dtypes=DTYPES, + traversal="1d", +) +``` + +The shared registrars derive the fallback priority/ID as `0/0` and the +preferred 1D priority/ID as `10/1`. Bespoke family registrars should use +`traversal_metadata(...)` to retain this policy; multi-form families pass their +fallback ID and form count so IDs remain unique. Candidate ID expresses +identity, not preference. Every ID for the same operation must remain unique. +The 1D registrar constraint must name every tile operand that participates in +or constrains the traversal, including temporary TileOp operands. + +The first production users of this pattern are the ordinary unary operations +`tabs`, `texp`, `tneg`, `tnot`, `trelu`, `trsqrt`, and `tsqrt`. Their original +candidate names and ID 0 remain the 2D fallbacks; their `_1d` candidates use +ID 1 and higher priority. + +The ordinary Tile-Tile operations `tadd`, `tand`, `tmax`, `tmin`, `tmul`, +`tor`, `tshl`, `tshr`, and `tsub` follow the same pattern through +`register_binary`. Even when an operation previously had an equivalent local +row-wise body, as `tmin` did, keep the ordinary computation as the vector +callback and let the shared registrar own both traversal forms. + +The ordinary Tile-Scalar operations use `register_scalar_binary` with the same +ID-0 fallback and preferred ID-1 pattern. Flattened legality names only the +traversed `src` and `dst` tiles; a scalar operand has no layout or continuity +requirement. Preserve instruction-specific scalar forms: bitwise and +tile-minus-scalar operations request vector broadcast, while scalar shift +operations keep their `i16` scalar dtype. `tsubs` is expressed as broadcast +plus `vsub` through the shared registrar rather than owning separate loops. + +Specialized Tile-Scalar algorithms remain in their family modules and call the +shared traversal emitters. `tdivs` retains its precision handling and distinct +tile-scalar/scalar-tile call forms; `_remainder.py` retains scalar remainder +math; and `tlrelu.py` retains slope coercion. Temporary scalar forms must name +the temporary in `require_elementwise_1d`, even when the generated body does +not access it. + +For callable forms distinguished by positional operand kind or order, test +selection through daemon metadata or a compiler lit test. Daemon requests +retain the original positional MLIR operand sequence. A direct registry call +uses a name-keyed mapping that is rebound to each candidate and therefore +cannot by itself prove which positional overload is legal. + +Scalar-fill operations use `register_scalar_fill`. Their 1D constraint names +only the destination tile because the scalar has no physical layout. Preserve +the original dtype signatures and let the shared emitter own flattened element +count, mask generation, tail handling, and vector stores. + +For bespoke computation, use the lower-level `emit_elementwise_1d` and +`emit_elementwise_2d` chunk callbacks or the family emitters. The flattened +form supplies one linear element offset; do not convert it back into row and +column indices. Predicate and dtype-width-changing conversion operations need +their additional representation-specific legality before using the 1D form. + +Keep bespoke algorithms in their operation modules. For example, `tlog` owns +its high-precision subnormal compensation and `trecip` owns its `1 / src` +calculation; both call shared unary traversal emitters. Do not move +operation-specific constants, precision algorithms, temporary calculations, or +instruction sequences into `_elementwise.py`. + +The same boundary applies to specialized binary families. `tdiv.py` owns its +precision-dependent division callback, and `_remainder.py` owns fmod/remainder +instruction sequences and their family registrars. These modules may call +shared binary traversal emitters, but their algorithms do not belong in +`_elementwise.py`. + +Temporary-operand Tile-Tile forms follow the same registration pattern, but +the 1D constraint must include the temporary explicitly. For example, +`tprelu`, `trem`, and `txor` call their family registrar with `has_tmp=True`; +the registrar passes `src0`, `src1`, `tmp`, and `dst` to +`require_elementwise_1d`. Keep the temporary in that proof even when the +current generated helper does not read it. A non-contiguous or insufficiently +described temporary must independently disqualify the 1D candidate. + +The issue-scope acceptance catalog lives in +`ptodsl/tests/test_tilelib_elementwise.py` as +`ELEMENTWISE_SCOPE_BY_FAMILY`. It is the authoritative checklist for the +unary, Tile-Tile, Tile-Scalar, compare, select, conversion, and scalar-fill +operations covered by shared 1D/2D selection. The catalog test requires unique +candidate IDs, a general 2D fallback, and higher-ranked 1D coverage for every +listed operation. There are currently no operation-level 1D exceptions. If an +instruction representation makes all flattened forms illegal in the future, +add a non-empty reviewed reason to `ELEMENTWISE_1D_EXCEPTIONS` and a focused +legality regression; do not remove the operation from the catalog. + ## View And Valid-Shape Rules Do not assume the logical valid shape is the same as the physical view shape. diff --git a/ptodsl/ptodsl/_types.py b/ptodsl/ptodsl/_types.py index bc983994cf..96aea46f38 100644 --- a/ptodsl/ptodsl/_types.py +++ b/ptodsl/ptodsl/_types.py @@ -592,7 +592,8 @@ def tile_buf_type(shape, dtype, valid_shape=None, *, address_space: str = "ub", slayout: str = "NoneBox", fractal_size: int = 512, - pad: str = "Null") -> Type: + pad: str = "Null", + compact_mode: str | int = "Null") -> Type: """ Construct a ``!pto.tile_buf<…>`` type via the Python bindings. @@ -615,6 +616,7 @@ def tile_buf_type(shape, dtype, valid_shape=None, *, _pto.SLayoutAttr.get(getattr(_pto.SLayout, slayout)), fractal_size, _pto.PadValueAttr.get(getattr(_pto.PadValue, pad)), + compact_mode=_compact_mode_token(compact_mode), ) if valid_shape is None and cfg is None: return _pto.TileBufType.get(shape, elem, space_attr) @@ -625,6 +627,30 @@ def tile_buf_type(shape, dtype, valid_shape=None, *, return _pto.TileBufType.get(shape, elem, space_attr, valid_shape=valid_shape, config=cfg) +def _normalize_compact_mode(value: str | int): + if isinstance(value, int): + return value + aliases = { + "null": "Null", + "normal": "Normal", + "row_plus_one": "RowPlusOne", + "Null": "Null", + "Normal": "Normal", + "RowPlusOne": "RowPlusOne", + } + return aliases.get(str(value), str(value)) + + +def _compact_mode_token(value: str | int): + token = _normalize_compact_mode(value) + if isinstance(token, int): + return token + try: + return getattr(_pto.CompactMode, token) + except AttributeError as exc: + raise ValueError(f"Unknown compact_mode {value!r}") from exc + + def tensor_view_type(rank: int, elem) -> Type: """``!pto.tensor_view`` with *rank* all-dynamic dims.""" return _pto.TensorViewType.get(rank, _ensure_tensor_storage_dtype(elem, context="pto.tensor_view_type(...)")) diff --git a/ptodsl/ptodsl/tilelib/__init__.py b/ptodsl/ptodsl/tilelib/__init__.py index 5dd9a08157..017fe89cd6 100644 --- a/ptodsl/ptodsl/tilelib/__init__.py +++ b/ptodsl/ptodsl/tilelib/__init__.py @@ -25,6 +25,10 @@ check_s_layout, check_type, require_contiguous, + require_conversion_1d, + require_elementwise_1d, + require_predicate_compare_1d, + require_predicate_select_1d, require_same_valid_shape, require_valid_rows, ) @@ -85,6 +89,10 @@ "check_layout", "check_s_layout", "require_contiguous", + "require_conversion_1d", + "require_elementwise_1d", + "require_predicate_compare_1d", + "require_predicate_select_1d", "require_same_valid_shape", "require_valid_rows", "f32", diff --git a/ptodsl/ptodsl/tilelib/_selection.py b/ptodsl/ptodsl/tilelib/_selection.py index d34ebe51f7..3dc8636e50 100644 --- a/ptodsl/ptodsl/tilelib/_selection.py +++ b/ptodsl/ptodsl/tilelib/_selection.py @@ -11,8 +11,8 @@ from . import constraints as _constraints from . import registry as _registry +from ._template_package import load_template from .metadata import ScalarSpec, ScalarType, TileSpec, VectorSpec, ViewSpec -from TileOps import load_template def _build_tile_specs(descriptor, operand_specs: list) -> dict: @@ -94,6 +94,9 @@ def _build_tile_specs(descriptor, operand_specs: list) -> dict: ) from exc valid_shape = spec.get("valid_shape") + s_fractal_size = config.get("s_fractal_size", 512) + if s_fractal_size == 0: + s_fractal_size = 512 specs[name] = TileSpec( shape=shape, dtype=dtype, @@ -101,7 +104,9 @@ def _build_tile_specs(descriptor, operand_specs: list) -> dict: valid_shape=tuple(valid_shape) if valid_shape is not None else None, b_layout=config.get("b_layout", "row_major"), s_layout=config.get("s_layout", "none_box"), + s_fractal_size=s_fractal_size, pad_value=spec.get("pad_value", config.get("pad_value", "Null")), + compact_mode=config.get("compact_mode", "null"), ) return specs @@ -207,10 +212,29 @@ def _legal_candidate_specs( f"no legal template for op={op!r} target={target!r}; {reasons}" ) - legal.sort(key=lambda pair: pair[0].metadata.priority, reverse=True) + legal.sort(key=lambda pair: _registry.candidate_sort_key(pair[0])) return legal +def _require_unambiguous_top_candidate(target: str, op: str, legal: list) -> None: + if len(legal) < 2: + return + + top_priority = legal[0][0].metadata.priority + winners = [ + descriptor + for descriptor, _ in legal + if descriptor.metadata.priority == top_priority + ] + if len(winners) > 1: + names = ", ".join(descriptor.name for descriptor in winners) + raise _registry.AmbiguousTemplate( + f"multiple templates tie at priority {top_priority} for op={op!r} " + f"target={target!r}: {names}; assign distinct priorities or make " + "their constraints mutually exclusive" + ) + + def _select_descriptor_and_specs( target: str, op: str, @@ -232,18 +256,7 @@ def _select_descriptor_and_specs( if len(legal) == 1: return legal[0] - top_priority = legal[0][0].metadata.priority - winners = [ - (descriptor, specs) - for descriptor, specs in legal - if descriptor.metadata.priority == top_priority - ] - if len(winners) > 1: - names = ", ".join(descriptor.name for descriptor, _ in winners) - raise _registry.AmbiguousTemplate( - f"multiple templates tie at priority {top_priority} for op={op!r} " - f"target={target!r}: {names}" - ) + _require_unambiguous_top_candidate(target, op, legal) return legal[0] @@ -253,13 +266,14 @@ def metadata_request( operand_specs: list, context_attrs: dict | None = None, ) -> dict: - """Return every legal candidate and its selection metadata.""" + """Return every legal candidate in deterministic selection order.""" legal = _legal_candidate_specs(target, op, operand_specs, context_attrs) + _require_unambiguous_top_candidate(target, op, legal) return { "target": target, "op": op, - "candidates": { - descriptor.name: _metadata_for_descriptor( + "candidates": [ + _metadata_for_descriptor( descriptor, { **_constraints.build_context(specs, target, op), @@ -267,7 +281,7 @@ def metadata_request( }, ) for descriptor, specs in legal - }, + ], } diff --git a/ptodsl/ptodsl/tilelib/_template_package.py b/ptodsl/ptodsl/tilelib/_template_package.py new file mode 100644 index 0000000000..2d49e735dc --- /dev/null +++ b/ptodsl/ptodsl/tilelib/_template_package.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""Resolve the canonical TileOps package in source and packaged environments.""" + +from __future__ import annotations + +from functools import lru_cache +from importlib import import_module + + +_SOURCE_PACKAGE = "TileOps" +_PACKAGED_PACKAGE = "ptoas._runtime.share.ptoas.TileOps" + + +def _is_missing_package(error: ModuleNotFoundError, package: str) -> bool: + missing = error.name + return bool(missing) and (missing == package or package.startswith(f"{missing}.")) + + +@lru_cache(maxsize=1) +def tileops_package(): + """Return TileOps from a source root or the bundled PTOAS resources.""" + + try: + return import_module(_SOURCE_PACKAGE) + except ModuleNotFoundError as source_error: + if not _is_missing_package(source_error, _SOURCE_PACKAGE): + raise + + try: + return import_module(_PACKAGED_PACKAGE) + except ModuleNotFoundError as packaged_error: + if not _is_missing_package(packaged_error, _PACKAGED_PACKAGE): + raise + raise ModuleNotFoundError( + "unable to locate TileOps in either the Python path or the " + "packaged PTOAS runtime resources", + name=_SOURCE_PACKAGE, + ) from packaged_error + + +def load_template(op: str, target: str) -> bool: + """Load a template through the resolved canonical TileOps package.""" + + return tileops_package().load_template(op, target) + + +__all__ = ["load_template", "tileops_package"] diff --git a/ptodsl/ptodsl/tilelib/constraints.py b/ptodsl/ptodsl/tilelib/constraints.py index 981e7a5459..ec67d7630f 100644 --- a/ptodsl/ptodsl/tilelib/constraints.py +++ b/ptodsl/ptodsl/tilelib/constraints.py @@ -28,6 +28,7 @@ from dataclasses import dataclass from enum import Enum +from .._types import _normalize_compact_mode from .metadata import ScalarSpec, VectorSpec, ViewSpec @@ -52,11 +53,12 @@ class CandidateLegality: @dataclass(frozen=True) class _ConfigView: - """The ``{name}_config`` object a constraint sees (``.b_layout`` / ``.s_layout`` strings, - which compare equal to the BLayout/SLayout str-enums).""" + """The ``{name}_config`` object exposed to constraint predicates.""" b_layout: str s_layout: str + s_fractal_size: int | None + compact_mode: str | int | None def build_context(tile_specs: dict, target: str, op: str) -> dict: @@ -68,9 +70,12 @@ def build_context(tile_specs: dict, target: str, op: str) -> dict: operand_rows = [] operand_cols = [] operand_sizes = [] + operand_valid_rows = [] operand_valid_cols = [] operand_b_layouts = [] operand_s_layouts = [] + operand_s_fractal_sizes = [] + operand_compact_modes = [] for name, spec in tile_specs.items(): dtype = spec.dtype.name operand_dtypes.append(dtype) @@ -122,33 +127,53 @@ def build_context(tile_specs: dict, target: str, op: str) -> dict: memory_space = getattr(spec, "memory_space", "ub") b_layout = getattr(spec, "b_layout", "row_major") s_layout = getattr(spec, "s_layout", "none_box") + s_fractal_size = getattr(spec, "s_fractal_size", None) + compact_mode = getattr(spec, "compact_mode", None) operand_memory_spaces.append(memory_space) operand_sizes.append(_shape_size(shape)) operand_b_layouts.append(b_layout) operand_s_layouts.append(s_layout) + operand_s_fractal_sizes.append(s_fractal_size) + operand_compact_modes.append(compact_mode) context[f"{name}_kind"] = "tile" context[f"{name}_shape"] = shape context[f"{name}_valid_shape"] = valid context[f"{name}_memory_space"] = memory_space + context[f"{name}_s_fractal_size"] = s_fractal_size + context[f"{name}_compact_mode"] = compact_mode context[f"{name}_config"] = _ConfigView( b_layout=b_layout, s_layout=s_layout, + s_fractal_size=s_fractal_size, + compact_mode=compact_mode, ) if len(shape) == 2: context[f"{name}_rows"], context[f"{name}_cols"] = shape - context[f"{name}_valid_rows"], context[f"{name}_valid_cols"] = valid + if len(valid) == 2: + ( + context[f"{name}_valid_rows"], + context[f"{name}_valid_cols"], + ) = valid operand_rows.append(shape[0]) operand_cols.append(shape[1]) - operand_valid_cols.append(valid[1]) + if len(valid) == 2: + operand_valid_rows.append(valid[0]) + operand_valid_cols.append(valid[1]) + else: + operand_valid_rows.append(None) + operand_valid_cols.append(None) context["operand_dtypes"] = tuple(operand_dtypes) context["operand_kinds"] = tuple(operand_kinds) context["operand_memory_spaces"] = tuple(operand_memory_spaces) context["operand_rows"] = tuple(operand_rows) context["operand_cols"] = tuple(operand_cols) context["operand_sizes"] = tuple(operand_sizes) + context["operand_valid_rows"] = tuple(operand_valid_rows) context["operand_valid_cols"] = tuple(operand_valid_cols) context["operand_b_layouts"] = tuple(operand_b_layouts) context["operand_s_layouts"] = tuple(operand_s_layouts) + context["operand_s_fractal_sizes"] = tuple(operand_s_fractal_sizes) + context["operand_compact_modes"] = tuple(operand_compact_modes) return context @@ -304,13 +329,416 @@ def require_contiguous(required=True): def _require_contiguous(operand_rows, operand_cols, operand_valid_cols, **_): if not required: return True - full_cols = all(valid == cols for valid, cols in zip(operand_valid_cols, operand_cols)) + if ( + len(operand_valid_cols) != len(operand_cols) + or None in operand_valid_cols + ): + return False + full_cols = all( + valid == cols + for valid, cols in zip(operand_valid_cols, operand_cols) + ) single_row = all(rows == 1 for rows in operand_rows) return full_cols or single_row return _require_contiguous +def require_elementwise_1d(*operand_names, memory_spaces=("ub", "vec")): + """Require ordinary element-wise operands to describe one flat range. + + The rule is intentionally limited to rank-2 local tiles with a row-major, + unboxed, gap-free physical layout. All named tiles must have the same + static logical valid shape. A range is flattenable when every tile uses + its full physical column axis, or when the logical range occupies only the + first row. Predicate and conversion representations need family-specific + constraints in addition to, or instead of, this ordinary rule. + """ + + allowed_memory_spaces = frozenset(memory_spaces) + + def _require_elementwise_1d(**context): + if not operand_names: + return False + + shapes = [] + valid_shapes = [] + for name in operand_names: + if not _is_flat_local_tile(context, name, allowed_memory_spaces): + return False + + shape = context.get(f"{name}_shape") + valid_shape = context.get(f"{name}_valid_shape") + shapes.append(shape) + valid_shapes.append(valid_shape) + + if any(valid_shape != valid_shapes[0] for valid_shape in valid_shapes[1:]): + return False + + full_columns = all( + valid_shape[1] == shape[1] + for shape, valid_shape in zip(shapes, valid_shapes) + ) + single_logical_row = valid_shapes[0][0] == 1 + return full_columns or single_logical_row + + return _require_elementwise_1d + + +def require_conversion_1d( + source_operand="src", + destination_operand="dst", + *, + source_elements_per_destination=1, + memory_spaces=("ub", "vec"), +): + """Require two typed conversion streams to be independently flattenable. + + Conversion source and destination tiles need not have the same element + width. Their typed pointers still advance over one common logical range, + provided each tile is gap-free and the valid shapes obey the conversion's + element-count relationship. ``source_elements_per_destination`` models + packed forms such as A5 BF16-to-FP4, where one destination storage element + represents two source elements. + + Multi-row ranges must fill the physical column axis of both tiles. A + single logical row is also legal because neither stream crosses a row + boundary. Unknown layout or compact-mode metadata rejects the candidate. + """ + + allowed_memory_spaces = frozenset(memory_spaces) + ratio = source_elements_per_destination + + def _require_conversion_1d(**context): + if not isinstance(ratio, int) or ratio <= 0: + return False + + operands = (source_operand, destination_operand) + shapes = [] + valid_shapes = [] + for name in operands: + if not _is_flat_local_tile(context, name, allowed_memory_spaces): + return False + + shape = context.get(f"{name}_shape") + valid_shape = context.get(f"{name}_valid_shape") + shapes.append(shape) + valid_shapes.append(valid_shape) + + src_shape, dst_shape = shapes + src_valid, dst_valid = valid_shapes + if src_shape[0] != dst_shape[0] or src_valid[0] != dst_valid[0]: + return False + if ( + src_shape[1] != dst_shape[1] * ratio + or src_valid[1] != dst_valid[1] * ratio + ): + return False + + full_columns = ( + src_valid[1] == src_shape[1] + and dst_valid[1] == dst_shape[1] + ) + return full_columns or dst_valid[0] == 1 + + return _require_conversion_1d + + +_PREDICATE_PACKING_LAYOUTS = { + # A5 stores one predicate bit per compared element. 32-bit comparisons + # combine two 64-lane masks before a 16-byte PK store; 16-bit comparisons + # use one 128-lane 16-byte PK store; and 8-bit comparisons use one + # 256-lane 32-byte NORM store. + "f32": (128, 16), + "i32": (128, 16), + "f16": (128, 16), + "i16": (128, 16), + "i8": (256, 32), + "ui8": (256, 32), +} + + +def require_predicate_compare_1d( + *data_operand_names, + predicate_operand="dst", + flattened_destination=False, + memory_spaces=("ub", "vec"), +): + """Require an A5 compare and its packed predicate output to be flattenable. + + Data operands must describe the same static contiguous logical range. The + predicate destination has a different representation: one bit per source + element, written in complete dtype-dependent predicate-store blocks. + + A single logical row is flattenable when its destination row has enough + physical bytes for the rounded store. Multiple rows normally require every + source row to end on a predicate-store boundary and the destination row + stride to equal the exact bytes produced for one row. + + ``flattened_destination`` models operations such as ``tcmps`` whose 1D + form writes the packed predicate into one continuous destination prefix. + It accepts either an exact packed row stride or a destination row with one + predicate container element per physical data element. The latter is an + explicit logical-range capacity contract, not a packed-byte requirement; + arbitrary intermediate predicate-row padding remains a 2D fallback. + """ + + allowed_memory_spaces = frozenset(memory_spaces) + + def _require_predicate_compare_1d(**context): + if not data_operand_names: + return False + + data_shapes = [] + data_valid_shapes = [] + for name in data_operand_names: + if not _is_flat_local_tile(context, name, allowed_memory_spaces): + return False + shape = context.get(f"{name}_shape") + valid_shape = context.get(f"{name}_valid_shape") + data_shapes.append(shape) + data_valid_shapes.append(valid_shape) + + if any( + valid_shape != data_valid_shapes[0] + for valid_shape in data_valid_shapes[1:] + ): + return False + + predicate_name = predicate_operand + if not _is_flat_local_tile( + context, + predicate_name, + allowed_memory_spaces, + ): + return False + + predicate_shape = context.get(f"{predicate_name}_shape") + predicate_valid_shape = context.get(f"{predicate_name}_valid_shape") + valid_rows, valid_cols = data_valid_shapes[0] + if predicate_valid_shape[0] != valid_rows: + return False + + dtype = context.get(f"{data_operand_names[0]}_dtype") + store_layout = _PREDICATE_PACKING_LAYOUTS.get(dtype) + if store_layout is None: + return False + elements_per_store, bytes_per_store = store_layout + + predicate_dtype = context.get(f"{predicate_name}_dtype") + predicate_bytewidth = _dtype_bytewidth(predicate_dtype) + if predicate_bytewidth is None: + return False + predicate_row_bytes = predicate_shape[1] * predicate_bytewidth + if predicate_row_bytes % 32 != 0: + return False + row_store_count = _ceil_div(valid_cols, elements_per_store) + required_row_bytes = row_store_count * bytes_per_store + if predicate_row_bytes < required_row_bytes: + return False + + single_logical_row = valid_rows == 1 + if single_logical_row: + required_data_row_elements = row_store_count * elements_per_store + return all( + shape[1] >= required_data_row_elements + for shape in data_shapes + ) + + data_rows_are_contiguous = all( + valid_shape[1] == shape[1] + for shape, valid_shape in zip(data_shapes, data_valid_shapes) + ) + required_logical_row_bytes = ( + max(shape[1] for shape in data_shapes) * predicate_bytewidth + ) + destination_holds_logical_range = ( + predicate_row_bytes >= required_logical_row_bytes + ) + if flattened_destination and destination_holds_logical_range: + total_elements = valid_rows * valid_cols + total_store_count = _ceil_div( + total_elements, + elements_per_store, + ) + required_total_bytes = total_store_count * bytes_per_store + predicate_total_bytes = ( + predicate_shape[0] * predicate_row_bytes + ) + return ( + data_rows_are_contiguous + and predicate_total_bytes >= required_total_bytes + ) + + rows_end_on_store_boundary = valid_cols % elements_per_store == 0 + predicate_rows_are_contiguous = predicate_row_bytes == required_row_bytes + return ( + data_rows_are_contiguous + and rows_end_on_store_boundary + and predicate_rows_are_contiguous + ) + + return _require_predicate_compare_1d + + +def require_predicate_select_1d( + predicate_operand, + *data_operand_names, + temporary_operand=None, + memory_spaces=("ub", "vec"), +): + """Require A5 packed-predicate select operands to be flattenable. + + Data operands must describe one static contiguous logical range. The mask + is a byte-addressed packed predicate whose row capacity and stride are + checked using the selected data dtype. A5 does not access the ABI + temporary, but its tile metadata must still be complete and supported. + """ + + allowed_memory_spaces = frozenset(memory_spaces) + + def _require_predicate_select_1d(**context): + if not data_operand_names: + return False + + data_shapes = [] + data_valid_shapes = [] + for name in data_operand_names: + if not _is_flat_local_tile(context, name, allowed_memory_spaces): + return False + data_shapes.append(context.get(f"{name}_shape")) + data_valid_shapes.append(context.get(f"{name}_valid_shape")) + + if any( + valid_shape != data_valid_shapes[0] + for valid_shape in data_valid_shapes[1:] + ): + return False + + if not _is_flat_local_tile( + context, + predicate_operand, + allowed_memory_spaces, + ): + return False + if temporary_operand is not None and not _is_flat_local_tile( + context, + temporary_operand, + allowed_memory_spaces, + ): + return False + + predicate_shape = context.get(f"{predicate_operand}_shape") + predicate_valid_shape = context.get( + f"{predicate_operand}_valid_shape" + ) + valid_rows, valid_cols = data_valid_shapes[0] + if predicate_valid_shape[0] != valid_rows: + return False + + dtype = context.get(f"{data_operand_names[0]}_dtype") + store_layout = _PREDICATE_PACKING_LAYOUTS.get(dtype) + if store_layout is None: + return False + elements_per_store, bytes_per_store = store_layout + + predicate_dtype = context.get(f"{predicate_operand}_dtype") + predicate_bytewidth = _dtype_bytewidth(predicate_dtype) + if predicate_bytewidth is None: + return False + predicate_row_bytes = predicate_shape[1] * predicate_bytewidth + if predicate_row_bytes % 32 != 0: + return False + + row_store_count = _ceil_div(valid_cols, elements_per_store) + required_predicate_row_bytes = row_store_count * bytes_per_store + if predicate_row_bytes < required_predicate_row_bytes: + return False + + single_logical_row = valid_rows == 1 + if single_logical_row: + required_data_row_elements = row_store_count * elements_per_store + return all( + shape[1] >= required_data_row_elements + for shape in data_shapes + ) + + data_rows_are_contiguous = all( + valid_shape[1] == shape[1] + for shape, valid_shape in zip(data_shapes, data_valid_shapes) + ) + rows_end_on_store_boundary = valid_cols % elements_per_store == 0 + predicate_rows_are_contiguous = ( + predicate_row_bytes == required_predicate_row_bytes + ) + return ( + data_rows_are_contiguous + and rows_end_on_store_boundary + and predicate_rows_are_contiguous + ) + + return _require_predicate_select_1d + + +def _is_flat_local_tile(context, name, allowed_memory_spaces) -> bool: + if context.get(f"{name}_kind") != "tile": + return False + shape = context.get(f"{name}_shape") + valid_shape = context.get(f"{name}_valid_shape") + if not _is_static_rank2_shape(shape) or not _is_static_rank2_shape( + valid_shape + ): + return False + if any(valid > physical for valid, physical in zip(valid_shape, shape)): + return False + if context.get(f"{name}_memory_space") not in allowed_memory_spaces: + return False + + config = context.get(f"{name}_config") + return ( + config is not None + and _enum_value(config.b_layout) == BLayout.ROW_MAJOR.value + and _enum_value(config.s_layout) == SLayout.NONE_BOX.value + and _has_gap_free_row_stride(config.compact_mode) + ) + + +def _dtype_bytewidth(dtype) -> int | None: + widths = { + "i8": 1, + "ui8": 1, + "i16": 2, + "ui16": 2, + "f16": 2, + "bf16": 2, + "i32": 4, + "ui32": 4, + "f32": 4, + } + return widths.get(dtype) + + +def _ceil_div(value, divisor): + return (value + divisor - 1) // divisor + + +def _is_static_rank2_shape(shape) -> bool: + return ( + isinstance(shape, tuple) + and len(shape) == 2 + and all(isinstance(dim, int) and dim > 0 for dim in shape) + ) + + +def _has_gap_free_row_stride(compact_mode) -> bool: + return _normalize_compact_mode(_enum_value(compact_mode)) in { + 0, + 1, + "Null", + "Normal", + } + + def passes(predicates, context: dict) -> bool: """Return True iff every predicate is satisfied for *context* (legality filter).""" for predicate in predicates: @@ -353,6 +781,10 @@ def passes(predicates, context: dict) -> bool: "evaluate_candidate", "passes", "require_contiguous", + "require_conversion_1d", + "require_elementwise_1d", + "require_predicate_compare_1d", + "require_predicate_select_1d", "require_same_valid_shape", "require_valid_rows", ] diff --git a/ptodsl/ptodsl/tilelib/metadata.py b/ptodsl/ptodsl/tilelib/metadata.py index 9d39e52f16..3828cb57f4 100644 --- a/ptodsl/ptodsl/tilelib/metadata.py +++ b/ptodsl/ptodsl/tilelib/metadata.py @@ -118,8 +118,9 @@ def _scalar_type_token(dtype: ScalarType) -> str: class TileSpec: """Concrete specialization of one tile operand. - ``valid_shape``/``b_layout``/``s_layout``/``memory_space`` are carried for both - constraint evaluation (selection) and the rendered entry ``tile_buf`` type. + Shape, valid shape, memory space, and the complete tile configuration are + carried for both constraint evaluation (selection) and the rendered entry + ``tile_buf`` type. """ shape: tuple @@ -128,7 +129,9 @@ class TileSpec: valid_shape: tuple | None = None b_layout: str = "row_major" s_layout: str = "none_box" + s_fractal_size: int = 512 pad_value: str = "Null" + compact_mode: str | int = "null" def __post_init__(self): if len(self.shape) != 2: @@ -146,8 +149,9 @@ def mlir_type(self): blayout=_layout_token(self.b_layout), address_space=self.memory_space, slayout=_layout_token(self.s_layout), - fractal_size=512, + fractal_size=self.s_fractal_size, pad=_pad_token(self.pad_value), + compact_mode=self.compact_mode, ) diff --git a/ptodsl/ptodsl/tilelib/registry.py b/ptodsl/ptodsl/tilelib/registry.py index c437a0a8ed..c4af0933e7 100644 --- a/ptodsl/ptodsl/tilelib/registry.py +++ b/ptodsl/ptodsl/tilelib/registry.py @@ -31,6 +31,17 @@ class AmbiguousTemplate(Exception): pass +def candidate_sort_key(descriptor): + """Return the deterministic legal-candidate reporting order. + + Priority is the only selection rank. Name is a stable tie ordering for + diagnostics and non-winning candidates; a top-priority tie is still an + ambiguity and is rejected by ``select``. + """ + + return (-descriptor.metadata.priority, descriptor.name) + + class TileTemplateRegistry: def __init__(self): self._descriptors: list = [] @@ -82,7 +93,7 @@ def legal_candidates(self, op: str, target: str, tile_specs: dict, f"no legal template for op={op!r} target={target!r}; {reasons}" ) - legal.sort(key=lambda d: d.metadata.priority, reverse=True) + legal.sort(key=candidate_sort_key) return legal def select(self, op: str, target: str, tile_specs: dict, @@ -106,7 +117,9 @@ def select(self, op: str, target: str, tile_specs: dict, if len(winners) > 1: names = ", ".join(d.name for d in winners) raise AmbiguousTemplate( - f"multiple templates tie at priority {top_priority} for op={op!r} target={target!r}: {names}" + f"multiple templates tie at priority {top_priority} for op={op!r} " + f"target={target!r}: {names}; assign distinct priorities or make " + "their constraints mutually exclusive" ) return legal[0] @@ -127,7 +140,7 @@ def _load_default_templates(op: str, target: str) -> None: # Import lazily to avoid a registry/templates import cycle during package # initialization. The loader is cached and registers descriptors as a # module-import side effect. - from TileOps import load_template + from ._template_package import load_template load_template(op, target) @@ -148,6 +161,7 @@ def select(op: str, target: str, tile_specs: dict, context_attrs: dict | None = "TileTemplateRegistry", "NoMatchingTemplate", "AmbiguousTemplate", + "candidate_sort_key", "default_registry", "legal_candidates", "register", diff --git a/ptodsl/ptodsl/tilelib/templates/__init__.py b/ptodsl/ptodsl/tilelib/templates/__init__.py index 6ed28847c2..cafa535f64 100644 --- a/ptodsl/ptodsl/tilelib/templates/__init__.py +++ b/ptodsl/ptodsl/tilelib/templates/__init__.py @@ -5,18 +5,15 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""Compatibility loader for PTODSL TileLib templates. - -The template sources live in the top-level ``TileOps`` package. -""" +"""Compatibility loader for the canonical source or packaged TileOps.""" from __future__ import annotations from pathlib import Path -from TileOps import load_template -import TileOps as _tileops +from .._template_package import load_template, tileops_package +_tileops = tileops_package() __path__ = [str(Path(_tileops.__file__).resolve().parent)] __all__ = ["load_template"] diff --git a/ptodsl/tests/test_ptoas_runtime.py b/ptodsl/tests/test_ptoas_runtime.py index d8a03c23d7..c202668e66 100644 --- a/ptodsl/tests/test_ptoas_runtime.py +++ b/ptodsl/tests/test_ptoas_runtime.py @@ -88,7 +88,9 @@ def counted_materialize(*args, **kwargs): ) self.assertEqual(result, 0) - self.assertEqual(calls, 1) + # The input has two identical 1D calls and one distinct 2D + # fallback, so only the duplicate 1D specialization is reused. + self.assertEqual(calls, 2) self.assertIn("pto.vadd", output.read_text(encoding="utf-8")) diff --git a/ptodsl/tests/test_tilelib_catalog.py b/ptodsl/tests/test_tilelib_catalog.py index 10f68a6780..4060c0fb6d 100644 --- a/ptodsl/tests/test_tilelib_catalog.py +++ b/ptodsl/tests/test_tilelib_catalog.py @@ -17,17 +17,17 @@ # op -> (template name, rendered op, parameter names, representative dtype[, candidate id]) CATALOG = { - "pto.tabs": ("template_tabs", "pto.vabs", ("src", "dst"), "f32"), + "pto.tabs": ("template_tabs_1d", "pto.vabs", ("src", "dst"), "f32"), "pto.tadd": ( - "template_tadd", + "template_tadd_1d", "pto.vadd", ("src0", "src1", "dst"), "f32", ), - "pto.tand": ("template_tand", "pto.vand", ("src0", "src1", "dst"), "i32"), - "pto.tands": ("template_tands", "pto.vand", ("src", "scalar", "dst"), "i32"), + "pto.tand": ("template_tand_1d", "pto.vand", ("src0", "src1", "dst"), "i32"), + "pto.tands": ("template_tands_1d", "pto.vand", ("src", "scalar", "dst"), "i32"), "pto.tcmp": ("template_tcmp", "pto.vcmp", ("src0", "src1", "dst"), "f32"), - "pto.tcmps": ("template_tcmps", "pto.vcmps", ("src", "scalar", "dst"), "f32"), + "pto.tcmps": ("template_tcmps_1d", "pto.vcmps", ("src", "scalar", "dst"), "f32"), "pto.tcolexpand": ("template_tcolexpand", "pto.vlds", ("src", "dst"), "f32"), "pto.tcolexpandadd": ("template_tcolexpandadd", "pto.vadd", ("src0", "src1", "dst"), "f32"), "pto.tcolexpanddiv": ("template_tcolexpanddiv", "pto.vdiv", ("src0", "src1", "dst"), "f32"), @@ -47,26 +47,26 @@ "pto.tcolmin": ("template_tcolmin", "pto.vmin", ("src", "dst"), "f32"), "pto.tcolprod": ("template_tcolprod", "pto.vmul", ("src", "dst"), "f32"), "pto.tcolsum": ("template_tcolsum", "pto.vadd", ("src", "dst"), "f32"), - "pto.texpands": ("template_texpands", "pto.vdup", ("scalar", "dst"), "f32"), + "pto.texpands": ("template_texpands_1d", "pto.vdup", ("scalar", "dst"), "f32"), "pto.textract": ("template_textract_vec2vec_nd", "pto.vlds", ("src", "index_row", "index_col", "dst"), "f32"), - "pto.tlrelu": ("template_tlrelu", "pto.vlrelu", ("src", "slope", "dst"), "f32"), - "pto.tlog": ("template_tlog", "pto.vln", ("src", "dst"), "f32"), - "pto.tdiv": ("template_tdiv", "pto.vdiv", ("src0", "src1", "dst"), "f32"), + "pto.tlrelu": ("template_tlrelu_1d", "pto.vlrelu", ("src", "slope", "dst"), "f32"), + "pto.tlog": ("template_tlog_1d", "pto.vln", ("src", "dst"), "f32"), + "pto.tdiv": ("template_tdiv_1d", "pto.vdiv", ("src0", "src1", "dst"), "f32"), "pto.tdivs": ( - "template_tdivs_tile_scalar", + "template_tdivs_tile_scalar_1d", "pto.vdiv", ("src", "scalar", "dst"), "f32", - "template_tdivs_tile_scalar", + "template_tdivs_tile_scalar_1d", ), - "pto.tcvt": ("template_tcvt_f32_to_i32", "pto.vcvt", ("src", "dst"), "f32"), + "pto.tcvt": ("template_tcvt_f32_to_i32_1d", "pto.vcvt", ("src", "dst"), "f32"), "pto.tconcat": ("template_tconcat", "pto.vsts", ("src0", "src1", "dst"), "f32"), # tdequant has i16/i8 variants; this entry covers the i16 representative # (the i8 path is exercised by test_tdequant_dtype_versions_render). "pto.tdequant": ("template_tdequant_i16", "pto.vmul", ("src", "scale", "offset", "dst"), "i16"), - "pto.texp": ("template_texp", "pto.vexp", ("src", "dst"), "f32"), - "pto.tfmod": ("template_tfmod", "pto.vtrc", ("src0", "src1", "dst"), "f32"), - "pto.tfmods": ("template_tfmods", "pto.vtrc", ("src", "scalar", "dst"), "f32"), + "pto.texp": ("template_texp_1d", "pto.vexp", ("src", "dst"), "f32"), + "pto.tfmod": ("template_tfmod_1d", "pto.vtrc", ("src0", "src1", "dst"), "f32"), + "pto.tfmods": ("template_tfmods_1d", "pto.vtrc", ("src", "scalar", "dst"), "f32"), "pto.tfillpad": ("template_tfillpad", "pto.vsts", ("src", "dst"), "f32"), "pto.tgemv": ("template_tgemv", "pto.mad", ("lhs", "rhs", "acc"), "f16"), "pto.tgemv.acc": ("template_tgemv_acc", "pto.mad_acc", ("acc_in", "lhs", "rhs", "dst"), "f16"), @@ -118,24 +118,24 @@ ("lhs", "lhs_scale", "rhs", "rhs_scale", "bias", "dst"), "f8e4m3", ), - "pto.tmax": ("template_tmax", "pto.vmax", ("src0", "src1", "dst"), "f32"), - "pto.tneg": ("template_tneg", "pto.vneg", ("src", "dst"), "f32"), - "pto.tmin": ("template_tmin", "pto.vmin", ("src0", "src1", "dst"), "f32"), + "pto.tmax": ("template_tmax_1d", "pto.vmax", ("src0", "src1", "dst"), "f32"), + "pto.tneg": ("template_tneg_1d", "pto.vneg", ("src", "dst"), "f32"), + "pto.tmin": ("template_tmin_1d", "pto.vmin", ("src0", "src1", "dst"), "f32"), "pto.tmov": ("template_tmov_basic", "pto.vsts", ("src", "dst"), "f32"), - "pto.tnot": ("template_tnot", "pto.vnot", ("src", "dst"), "i32"), - "pto.tor": ("template_tor", "pto.vor", ("src0", "src1", "dst"), "i32"), - "pto.tors": ("template_tors", "pto.vor", ("src", "scalar", "dst"), "i32"), + "pto.tnot": ("template_tnot_1d", "pto.vnot", ("src", "dst"), "i32"), + "pto.tor": ("template_tor_1d", "pto.vor", ("src0", "src1", "dst"), "i32"), + "pto.tors": ("template_tors_1d", "pto.vor", ("src", "scalar", "dst"), "i32"), "pto.tpartadd": ("template_tpartadd", "pto.vadd", ("src0", "src1", "dst"), "f32"), "pto.tpartmax": ("template_tpartmax", "pto.vmax", ("src0", "src1", "dst"), "f32"), "pto.tpartmin": ("template_tpartmin", "pto.vmin", ("src0", "src1", "dst"), "f32"), "pto.tpartmul": ("template_tpartmul", "pto.vmul", ("src0", "src1", "dst"), "f32"), - "pto.tprelu": ("template_tprelu", "pto.vprelu", ("src0", "src1", "tmp", "dst"), "f32"), + "pto.tprelu": ("template_tprelu_1d", "pto.vprelu", ("src0", "src1", "tmp", "dst"), "f32"), "pto.trandom": ("template_trandom", "pto.vmull", ("key0", "key1", "counter0", "counter1", "counter2", "counter3", "dst"), "ui32"), - "pto.trelu": ("template_trelu", "pto.vrelu", ("src", "dst"), "f32"), - "pto.trecip": ("template_trecip", "pto.vdiv", ("src", "dst"), "f32"), - "pto.trem": ("template_trem", "pto.vtrc", ("src0", "src1", "tmp", "dst"), "f32"), - "pto.trems": ("template_trems", "pto.vtrc", ("src", "scalar", "tmp", "dst"), "f32"), - "pto.trsqrt": ("template_trsqrt", "pto.vsqrt", ("src", "dst"), "f32"), + "pto.trelu": ("template_trelu_1d", "pto.vrelu", ("src", "dst"), "f32"), + "pto.trecip": ("template_trecip_1d", "pto.vdiv", ("src", "dst"), "f32"), + "pto.trem": ("template_trem_1d", "pto.vtrc", ("src0", "src1", "tmp", "dst"), "f32"), + "pto.trems": ("template_trems_1d", "pto.vtrc", ("src", "scalar", "tmp", "dst"), "f32"), + "pto.trsqrt": ("template_trsqrt_1d", "pto.vsqrt", ("src", "dst"), "f32"), "pto.trowargmax": ("template_trowargmax", "pto.vdintlv", ("src", "tmp", "dst"), "f32"), "pto.trowargmin": ("template_trowargmin", "pto.vdintlv", ("src", "tmp", "dst"), "f32"), "pto.trowexpand": ("template_trowexpand", "pto.vdup", ("src", "dst"), "f32"), @@ -157,10 +157,10 @@ "pto.trowsum": ("template_trowsum", "pto.vcadd", ("src", "tmp", "dst"), "f32"), "pto.tsel": ("template_tsel", "pto.vsel", ("mask", "src0", "src1", "tmp", "dst"), "f32"), "pto.tsels": ("template_tsels", "pto.vsel", ("mask", "src", "tmp", "scalar", "dst"), "f32"), - "pto.tshl": ("template_tshl", "pto.vshl", ("src0", "src1", "dst"), "i32"), - "pto.tshls": ("template_tshls", "pto.vshls", ("src", "scalar", "dst"), "i32"), - "pto.tshr": ("template_tshr", "pto.vshr", ("src0", "src1", "dst"), "i32"), - "pto.tshrs": ("template_tshrs", "pto.vshrs", ("src", "scalar", "dst"), "i32"), + "pto.tshl": ("template_tshl_1d", "pto.vshl", ("src0", "src1", "dst"), "i32"), + "pto.tshls": ("template_tshls_1d", "pto.vshls", ("src", "scalar", "dst"), "i32"), + "pto.tshr": ("template_tshr_1d", "pto.vshr", ("src0", "src1", "dst"), "i32"), + "pto.tshrs": ("template_tshrs_1d", "pto.vshrs", ("src", "scalar", "dst"), "i32"), "pto.tmrgsort": ("template_tmrgsort_multi_list2", "pto.vmrgsort4", ("src0", "src1", "tmp", "dst", "ex_vec"), "f32"), "pto.tsort32": ("template_tsort32", "pto.vbitsort", ("src", "idx", "dst"), "f32"), "pto.tstore": ( @@ -170,26 +170,26 @@ "f32", "template_tstore_nd", ), - "pto.tadds": ("template_tadds", "pto.vadds", ("src", "scalar", "dst"), "f32"), - "pto.tmaxs": ("template_tmaxs", "pto.vmaxs", ("src", "scalar", "dst"), "f32"), - "pto.tmins": ("template_tmins", "pto.vmins", ("src", "scalar", "dst"), "f32"), - "pto.tmuls": ("template_tmuls", "pto.vmuls", ("src", "scalar", "dst"), "f32"), + "pto.tadds": ("template_tadds_1d", "pto.vadds", ("src", "scalar", "dst"), "f32"), + "pto.tmaxs": ("template_tmaxs_1d", "pto.vmaxs", ("src", "scalar", "dst"), "f32"), + "pto.tmins": ("template_tmins_1d", "pto.vmins", ("src", "scalar", "dst"), "f32"), + "pto.tmuls": ("template_tmuls_1d", "pto.vmuls", ("src", "scalar", "dst"), "f32"), "pto.tmul": ( - "template_tmul", + "template_tmul_1d", "pto.vmul", ("src0", "src1", "dst"), "f32", ), "pto.txor": ( - "template_txor", + "template_txor_1d", "pto.vxor", ("src0", "src1", "tmp", "dst"), "i32", ), - "pto.txors": ("template_txors", "pto.vxor", ("src", "scalar", "tmp", "dst"), "i32"), - "pto.tsubs": ("template_tsubs", "pto.vsub", ("src", "scalar", "dst"), "f32"), - "pto.tsub": ("template_tsub", "pto.vsub", ("src0", "src1", "dst"), "f32"), - "pto.tsqrt": ("template_tsqrt", "pto.vsqrt", ("src", "dst"), "f32"), + "pto.txors": ("template_txors_1d", "pto.vxor", ("src", "scalar", "tmp", "dst"), "i32"), + "pto.tsubs": ("template_tsubs_1d", "pto.vsub", ("src", "scalar", "dst"), "f32"), + "pto.tsub": ("template_tsub_1d", "pto.vsub", ("src0", "src1", "dst"), "f32"), + "pto.tsqrt": ("template_tsqrt_1d", "pto.vsqrt", ("src", "dst"), "f32"), } CUBE_OPS = { @@ -257,6 +257,51 @@ OPS_WITHOUT_VECTOR_STORE = {"pto.tcmp", "pto.tcmps", "pto.tsort32"} OPS_WITHOUT_VECTOR_STORE = OPS_WITHOUT_VECTOR_STORE | {"pto.tload", "pto.tstore"} OPS_WITHOUT_VECTOR_STORE = OPS_WITHOUT_VECTOR_STORE | CUBE_OPS +OPS_WITHOUT_MEMREF_SUBVIEW = {"pto.tcmps", "pto.tsort32"} +OPS_WITHOUT_MEMREF_SUBVIEW = OPS_WITHOUT_MEMREF_SUBVIEW | {"pto.texpands", "pto.tdivs", "pto.tfillpad_inplace"} +OPS_WITHOUT_MEMREF_SUBVIEW = OPS_WITHOUT_MEMREF_SUBVIEW | {"pto.tload", "pto.tstore", "pto.tstore_fp", "pto.textract_fp"} +OPS_WITHOUT_MEMREF_SUBVIEW = OPS_WITHOUT_MEMREF_SUBVIEW | ROW_REDUCTIONS +OPS_WITHOUT_MEMREF_SUBVIEW = OPS_WITHOUT_MEMREF_SUBVIEW | ARG_COLUMN_REDUCTIONS +OPS_WITHOUT_MEMREF_SUBVIEW = OPS_WITHOUT_MEMREF_SUBVIEW | CUBE_OPS +OPS_WITHOUT_MEMREF_SUBVIEW = OPS_WITHOUT_MEMREF_SUBVIEW | { + "pto.tcvt", + "pto.tabs", + "pto.texp", + "pto.tneg", + "pto.tnot", + "pto.tlog", + "pto.trecip", + "pto.trelu", + "pto.trsqrt", + "pto.tsqrt", + "pto.tadd", + "pto.tand", + "pto.tmax", + "pto.tmin", + "pto.tmul", + "pto.tor", + "pto.tshl", + "pto.tshr", + "pto.tsub", + "pto.tdiv", + "pto.tfmod", + "pto.tprelu", + "pto.trem", + "pto.txor", + "pto.tadds", + "pto.tands", + "pto.tmaxs", + "pto.tmins", + "pto.tmuls", + "pto.tors", + "pto.tshls", + "pto.tshrs", + "pto.tsubs", + "pto.tfmods", + "pto.tlrelu", + "pto.trems", + "pto.txors", +} OPS_WITHOUT_LOOP = {"pto.tmrgsort"} OPS_WITHOUT_LOOP = OPS_WITHOUT_LOOP | CUBE_OPS OPS_ALLOWING_CASTPTR = {"pto.tsel", "pto.tsels"} @@ -835,7 +880,7 @@ def test_tcmps_vec_tiles_render_packed_mask_paths(self): ), } selected = select("pto.tcmps", "a5", specs) - self.assertEqual(selected.name, "template_tcmps") + self.assertEqual(selected.name, "template_tcmps_1d") mlir = selected.specialize(context_attrs={"cmp_mode": "lt"}, **specs).mlir_text() self.assertIn(expected_op, mlir) self.assertIn(expected_dist, mlir) @@ -871,7 +916,27 @@ def test_tmov_accepts_ui8_vec_tiles(self): self.assertEqual(selected.name, "template_tmov_basic") self.assertIn("pto.vsts", selected.specialize(**specs).mlir_text()) - def test_tcvt_additional_rowwise_versions_render(self): + def test_tfillpad_expanding_zero_pad_remains_zero(self): + specs = { + "src": TileSpec( + shape=(128, 128), + valid_shape=(128, 64), + dtype=ScalarType("f32"), + pad_value="Null", + ), + "dst": TileSpec( + shape=(128, 128), + valid_shape=(128, 128), + dtype=ScalarType("f32"), + pad_value="Zero", + ), + } + selected = select("pto.tfillpad", "a5", specs) + mlir = selected.specialize(**specs).mlir_text() + self.assertIn("arith.constant 0.000000e+00 : f32", mlir) + self.assertNotIn("arith.constant -1.000000e+00 : f32", mlir) + + def test_tcvt_contiguous_versions_select_flattened_candidates(self): signatures = { ("i32", "f32"): "template_tcvt_i32_to_f32", ("i16", "f16"): "template_tcvt_i16_to_f16", @@ -918,9 +983,12 @@ def test_tcvt_additional_rowwise_versions_render(self): "dst": TileSpec(shape=(8, 64), dtype=ScalarType(dst_dtype)), } selected = select("pto.tcvt", "a5", specs) - self.assertEqual(selected.name, expected_name) + self.assertEqual(selected.name, f"{expected_name}_1d") + self.assertEqual(selected.metadata.loop_depth, 1) expected_op = "pto.vtrc" if expected_name == "template_tcvt_f32_to_f32" else "pto.vcvt" - self.assertIn(expected_op, selected.specialize(**specs).mlir_text()) + mlir = selected.specialize(**specs).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn(expected_op, mlir) def test_tcvt_bf16_to_fp4_versions_render(self): for dst_dtype in ("f4e1m2x2", "f4e2m1x2"): @@ -930,8 +998,166 @@ def test_tcvt_bf16_to_fp4_versions_render(self): "dst": TileSpec(shape=(8, 64), dtype=ScalarType(dst_dtype)), } selected = select("pto.tcvt", "a5", specs) - self.assertEqual(selected.name, "template_tcvt_bf16_to_fp4") - self.assertIn("pto.vcvt", selected.specialize(**specs).mlir_text()) + self.assertEqual(selected.name, "template_tcvt_bf16_to_fp4_1d") + self.assertEqual(selected.metadata.loop_depth, 1) + mlir = selected.specialize(**specs).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn("pto.vcvt", mlir) + + def test_tcvt_32_to_ui8_store_mask_matches_source_chunk(self): + cases = ( + ( + "flattened", + (1, 128), + None, + "template_tcvt_i32_to_ui8_1d", + 1, + ), + ( + "rowwise", + (2, 128), + (2, 65), + "template_tcvt_i32_to_ui8", + 2, + ), + ) + for src_dtype in ("i32", "ui32"): + for label, shape, valid_shape, expected_name, loop_depth in cases: + with self.subTest(src_dtype=src_dtype, case=label): + specs = { + "src": TileSpec( + shape=shape, + valid_shape=valid_shape, + dtype=ScalarType(src_dtype), + ), + "dst": TileSpec( + shape=shape, + valid_shape=valid_shape, + dtype=ScalarType("ui8"), + ), + } + selected = select("pto.tcvt", "a5", specs) + mlir = selected.specialize(**specs).mlir_text() + + self.assertEqual( + selected.name, + expected_name.replace("i32", src_dtype), + ) + self.assertEqual(selected.metadata.loop_depth, loop_depth) + self.assertIn("pto.plt_b32", mlir) + self.assertIn("pto.pbitcast", mlir) + self.assertNotIn("pto.plt_b8", mlir) + + def test_tcvt_partial_multi_row_ranges_retain_2d_fallbacks(self): + cases = ( + ("f32", "i32", "template_tcvt_f32_to_i32"), + ("f16", "i16", "template_tcvt_f16_to_i16"), + ("i64", "f32", "template_tcvt_i64_to_f32"), + ) + for src_dtype, dst_dtype, expected_name in cases: + with self.subTest(signature=(src_dtype, dst_dtype)): + specs = { + "src": TileSpec( + shape=(8, 65), + valid_shape=(8, 63), + dtype=ScalarType(src_dtype), + ), + "dst": TileSpec( + shape=(8, 65), + valid_shape=(8, 63), + dtype=ScalarType(dst_dtype), + ), + } + selected = select("pto.tcvt", "a5", specs) + self.assertEqual(selected.name, expected_name) + self.assertEqual(selected.metadata.loop_depth, 2) + self.assertEqual( + selected.specialize(**specs).mlir_text().count("scf.for"), + 2, + ) + + fp4_specs = { + "src": TileSpec( + shape=(8, 130), + valid_shape=(8, 126), + dtype=ScalarType("bf16"), + ), + "dst": TileSpec( + shape=(8, 65), + valid_shape=(8, 63), + dtype=ScalarType("f4e1m2x2"), + ), + } + selected = select("pto.tcvt", "a5", fp4_specs) + self.assertEqual(selected.name, "template_tcvt_bf16_to_fp4") + self.assertEqual(selected.metadata.loop_depth, 2) + + def test_tcvt_single_row_and_stride_gap_selection(self): + single_row = { + "src": TileSpec( + shape=(4, 65), + valid_shape=(1, 63), + dtype=ScalarType("f32"), + ), + "dst": TileSpec( + shape=(4, 65), + valid_shape=(1, 63), + dtype=ScalarType("i32"), + ), + } + stride_gap = { + "src": TileSpec( + shape=(4, 65), + dtype=ScalarType("f32"), + compact_mode="row_plus_one", + ), + "dst": TileSpec( + shape=(4, 65), + dtype=ScalarType("i32"), + ), + } + + self.assertEqual( + select("pto.tcvt", "a5", single_row).name, + "template_tcvt_f32_to_i32_1d", + ) + self.assertEqual( + select("pto.tcvt", "a5", stride_gap).name, + "template_tcvt_f32_to_i32", + ) + + def test_tcvt_catalog_has_one_1d_pair_for_every_existing_candidate(self): + specs = { + "src": TileSpec(shape=(8, 64), dtype=ScalarType("f32")), + "dst": TileSpec(shape=(8, 64), dtype=ScalarType("i32")), + } + select("pto.tcvt", "a5", specs) + candidates = [ + descriptor + for descriptor in tilelib.default_registry().lookup( + "pto.tcvt", + "a5", + ) + ] + by_id = {descriptor.metadata.id: descriptor for descriptor in candidates} + + self.assertEqual(len(candidates), 76) + self.assertEqual(set(by_id), set(range(76))) + for fallback_id in range(38): + with self.subTest(fallback_id=fallback_id): + fallback = by_id[fallback_id] + flattened = by_id[38 + fallback_id] + self.assertEqual(fallback.metadata.loop_depth, 2) + self.assertEqual(fallback.metadata.priority, 0) + self.assertEqual(flattened.metadata.loop_depth, 1) + self.assertEqual(flattened.metadata.priority, 10) + self.assertEqual(flattened.name, f"{fallback.name}_1d") + + def test_tcvt_templates_use_native_python_loop_syntax(self): + from ptodsl.tilelib.templates.a5 import tcvt + + source = Path(tcvt.__file__).read_text(encoding="utf-8") + self.assertNotIn("pto.for_(", source) def test_tcolexpanddiv_i32_uses_float_divide_path(self): specs = { diff --git a/ptodsl/tests/test_tilelib_constraints.py b/ptodsl/tests/test_tilelib_constraints.py index d68ca8f00a..516e4983c7 100644 --- a/ptodsl/tests/test_tilelib_constraints.py +++ b/ptodsl/tests/test_tilelib_constraints.py @@ -14,7 +14,14 @@ import unittest from ptodsl.tilelib import ScalarType, TileSpec, VectorSpec, ViewSpec, select -from ptodsl.tilelib.constraints import build_context +from ptodsl.tilelib.constraints import ( + build_context, + passes, + require_conversion_1d, + require_elementwise_1d, + require_predicate_compare_1d, + require_predicate_select_1d, +) from ptodsl.tilelib.registry import NoMatchingTemplate F32 = ScalarType("f32") @@ -27,6 +34,72 @@ def _specs(*, dst_valid=(1, 64), dst_blayout="row_major", dst_slayout="none_box" return {"src": src, "dst": dst} +def _tile( + *, + shape=(8, 64), + valid_shape=None, + memory_space="ub", + b_layout="row_major", + s_layout="none_box", + compact_mode="null", +): + return TileSpec( + shape=shape, + dtype=F32, + valid_shape=valid_shape if valid_shape is not None else shape, + memory_space=memory_space, + b_layout=b_layout, + s_layout=s_layout, + compact_mode=compact_mode, + ) + + +def _ordinary_elementwise_1d(specs, *operand_names): + context = build_context(specs, "a5", "pto.example") + return passes((require_elementwise_1d(*operand_names),), context) + + +def _conversion_1d(specs, *, source_elements_per_destination=1): + context = build_context(specs, "a5", "pto.tcvt") + return passes( + ( + require_conversion_1d( + source_elements_per_destination=( + source_elements_per_destination + ), + ), + ), + context, + ) + + +def _predicate_compare_1d(specs, *data_operand_names): + context = build_context(specs, "a5", "pto.example") + return passes( + ( + require_predicate_compare_1d( + *data_operand_names, + predicate_operand="dst", + ), + ), + context, + ) + + +def _predicate_select_1d(specs, *data_operand_names): + context = build_context(specs, "a5", "pto.example") + return passes( + ( + require_predicate_select_1d( + "mask", + *data_operand_names, + temporary_operand="tmp", + ), + ), + context, + ) + + class TileLibConstraintTest(unittest.TestCase): def test_legal_colmax_selected(self): chosen = select("pto.tcolmax", "a5", _specs()) @@ -76,6 +149,317 @@ def test_context_tracks_view_and_vector_operands(self): self.assertEqual(context["aux_shape"], (4,)) self.assertEqual(context["aux_size"], 4) + def test_elementwise_1d_accepts_contiguous_multi_row_tiles(self): + specs = { + "src0": _tile(), + "src1": _tile(memory_space="vec"), + "dst": _tile(), + } + + self.assertTrue( + _ordinary_elementwise_1d(specs, "src0", "src1", "dst") + ) + + def test_elementwise_1d_accepts_a_contiguous_single_logical_row(self): + specs = { + "src": _tile(valid_shape=(1, 31)), + "dst": _tile(valid_shape=(1, 31)), + } + + self.assertTrue(_ordinary_elementwise_1d(specs, "src", "dst")) + + def test_elementwise_1d_rejects_partial_columns_across_rows(self): + specs = { + "src": _tile(valid_shape=(4, 63)), + "dst": _tile(valid_shape=(4, 63)), + } + + self.assertFalse(_ordinary_elementwise_1d(specs, "src", "dst")) + + def test_elementwise_1d_rejects_mismatched_logical_ranges(self): + specs = { + "src": _tile(valid_shape=(8, 64)), + "dst": _tile(valid_shape=(7, 64)), + } + + self.assertFalse(_ordinary_elementwise_1d(specs, "src", "dst")) + + def test_elementwise_1d_rejects_unknown_or_unsupported_metadata(self): + cases = { + "dynamic valid shape": { + "src": _tile(valid_shape=(None, 64)), + "dst": _tile(valid_shape=(None, 64)), + }, + "non-local memory": { + "src": _tile(memory_space="gm"), + "dst": _tile(), + }, + "column-major block layout": { + "src": _tile(b_layout="col_major"), + "dst": _tile(), + }, + "boxed sub-layout": { + "src": _tile(s_layout="row_major"), + "dst": _tile(), + }, + "stride gap": { + "src": _tile(compact_mode=2), + "dst": _tile(), + }, + "unknown compact mode": { + "src": _tile(compact_mode=None), + "dst": _tile(), + }, + } + + for label, specs in cases.items(): + with self.subTest(label=label): + self.assertFalse( + _ordinary_elementwise_1d(specs, "src", "dst") + ) + + def test_elementwise_1d_includes_temporary_tiles(self): + specs = { + "src": _tile(), + "tmp": _tile(shape=(8, 65), valid_shape=(8, 64)), + "dst": _tile(), + } + + self.assertTrue(_ordinary_elementwise_1d(specs, "src", "dst")) + self.assertFalse( + _ordinary_elementwise_1d(specs, "src", "tmp", "dst") + ) + + def test_conversion_1d_accepts_typed_contiguous_streams(self): + specs = { + "src": _tile(shape=(8, 65)), + "dst": TileSpec( + shape=(8, 65), + dtype=ScalarType("i16"), + valid_shape=(8, 65), + compact_mode="normal", + ), + } + self.assertTrue(_conversion_1d(specs)) + + def test_conversion_1d_accepts_a_partial_single_logical_row(self): + specs = { + "src": _tile(shape=(4, 65), valid_shape=(1, 63)), + "dst": TileSpec( + shape=(4, 65), + dtype=ScalarType("i16"), + valid_shape=(1, 63), + ), + } + self.assertTrue(_conversion_1d(specs)) + + def test_conversion_1d_rejects_partial_multi_row_or_stride_gap(self): + partial = { + "src": _tile(shape=(4, 65), valid_shape=(4, 63)), + "dst": TileSpec( + shape=(4, 65), + dtype=ScalarType("i16"), + valid_shape=(4, 63), + ), + } + stride_gap = { + "src": _tile(shape=(4, 65), compact_mode="row_plus_one"), + "dst": TileSpec( + shape=(4, 65), + dtype=ScalarType("i16"), + ), + } + self.assertFalse(_conversion_1d(partial)) + self.assertFalse(_conversion_1d(stride_gap)) + + def test_conversion_1d_models_bf16_to_fp4_packing(self): + legal = { + "src": TileSpec( + shape=(8, 130), + dtype=ScalarType("bf16"), + ), + "dst": TileSpec( + shape=(8, 65), + dtype=ScalarType("f4e1m2x2"), + ), + } + wrong_ratio = { + "src": TileSpec( + shape=(8, 129), + dtype=ScalarType("bf16"), + ), + "dst": legal["dst"], + } + self.assertTrue( + _conversion_1d( + legal, + source_elements_per_destination=2, + ) + ) + self.assertFalse( + _conversion_1d( + wrong_ratio, + source_elements_per_destination=2, + ) + ) + + def test_predicate_compare_1d_accepts_single_row_with_rounded_capacity(self): + specs = { + "src": _tile(shape=(4, 128), valid_shape=(1, 63)), + "dst": TileSpec( + shape=(4, 32), + valid_shape=(1, 8), + dtype=ScalarType("ui8"), + ), + } + + self.assertTrue(_predicate_compare_1d(specs, "src")) + + def test_predicate_compare_1d_accepts_dense_block_aligned_rows(self): + specs = { + "src0": _tile(shape=(4, 256)), + "src1": _tile(shape=(4, 256), memory_space="vec"), + "dst": TileSpec( + shape=(4, 32), + dtype=ScalarType("i8"), + ), + } + + self.assertTrue(_predicate_compare_1d(specs, "src0", "src1")) + + def test_predicate_compare_1d_rejects_row_tail_or_predicate_padding(self): + cases = { + "source row tail": { + "src": TileSpec( + shape=(4, 128), + dtype=ScalarType("i8"), + ), + "dst": TileSpec( + shape=(4, 32), + dtype=ScalarType("ui8"), + ), + }, + "predicate row padding": { + "src": _tile(shape=(4, 256)), + "dst": TileSpec( + shape=(4, 64), + dtype=ScalarType("ui8"), + ), + }, + "insufficient single-row capacity": { + "src": _tile(shape=(1, 65), valid_shape=(1, 63)), + "dst": TileSpec( + shape=(1, 32), + dtype=ScalarType("ui8"), + ), + }, + "predicate stride gap": { + "src": _tile(shape=(4, 256)), + "dst": TileSpec( + shape=(4, 32), + dtype=ScalarType("ui8"), + compact_mode=2, + ), + }, + } + + for label, specs in cases.items(): + with self.subTest(label=label): + self.assertFalse(_predicate_compare_1d(specs, "src")) + + def test_predicate_select_1d_accepts_single_row_and_dense_multi_row(self): + cases = { + "single row": { + "mask": TileSpec( + shape=(1, 32), + valid_shape=(1, 8), + dtype=ScalarType("i8"), + ), + "src": _tile(shape=(1, 128), valid_shape=(1, 63)), + "tmp": _tile(shape=(1, 32)), + "dst": _tile(shape=(1, 128), valid_shape=(1, 63)), + }, + "dense multi row with i32 mask container": { + "mask": TileSpec( + shape=(4, 8), + dtype=ScalarType("i32"), + ), + "src": _tile(shape=(4, 256)), + "tmp": _tile(shape=(1, 32)), + "dst": _tile(shape=(4, 256)), + }, + } + + for label, specs in cases.items(): + with self.subTest(label=label): + self.assertTrue(_predicate_select_1d(specs, "src", "dst")) + + def test_predicate_select_1d_rejects_ineligible_metadata(self): + common = { + "mask": TileSpec( + shape=(4, 32), + dtype=ScalarType("i8"), + ), + "src": _tile(shape=(4, 256)), + "tmp": _tile(shape=(1, 32)), + "dst": _tile(shape=(4, 256)), + } + cases = { + "partial data row": { + **common, + "src": _tile(shape=(4, 256), valid_shape=(4, 255)), + "dst": _tile(shape=(4, 256), valid_shape=(4, 255)), + }, + "predicate row padding": { + **common, + "mask": TileSpec( + shape=(4, 64), + dtype=ScalarType("i8"), + ), + }, + "predicate stride gap": { + **common, + "mask": TileSpec( + shape=(4, 32), + dtype=ScalarType("i8"), + compact_mode=2, + ), + }, + "temporary stride gap": { + **common, + "tmp": _tile(shape=(1, 32), compact_mode=2), + }, + "mismatched data range": { + **common, + "dst": _tile(shape=(4, 256), valid_shape=(3, 256)), + }, + } + + for label, specs in cases.items(): + with self.subTest(label=label): + self.assertFalse(_predicate_select_1d(specs, "src", "dst")) + + def test_context_carries_complete_tile_layout_metadata(self): + context = build_context( + { + "src": TileSpec( + shape=(8, 64), + dtype=F32, + s_fractal_size=32, + compact_mode=1, + ) + }, + "a5", + "pto.example", + ) + + self.assertEqual(context["src_s_fractal_size"], 32) + self.assertEqual(context["src_compact_mode"], 1) + self.assertEqual(context["src_config"].s_fractal_size, 32) + self.assertEqual(context["src_config"].compact_mode, 1) + self.assertEqual(context["operand_s_fractal_sizes"], (32,)) + self.assertEqual(context["operand_compact_modes"], (1,)) + if __name__ == "__main__": unittest.main() diff --git a/ptodsl/tests/test_tilelib_elementwise.py b/ptodsl/tests/test_tilelib_elementwise.py index 29a95eafdc..3604875702 100644 --- a/ptodsl/tests/test_tilelib_elementwise.py +++ b/ptodsl/tests/test_tilelib_elementwise.py @@ -5,12 +5,99 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""Phase-4 breadth test: each ported elementwise op selects + renders to the structured -abstraction, using the right vector op.""" +"""Element-wise template selection, traversal, and scope acceptance tests.""" +import ast import unittest +from pathlib import Path -from ptodsl.tilelib import ScalarType, TileSpec, select +from ptodsl import pto +import ptodsl.tilelib as tilelib +import ptodsl.tilelib.templates.a5._elementwise as elementwise +import ptodsl.tilelib.templates.a5.tsel as tsel_templates +import ptodsl.tilelib.templates.a5.tsels as tsels_templates +from ptodsl.tilelib import ( + ScalarSpec, + ScalarType, + TileSpec, + legal_candidates, + select, +) +from ptodsl.tilelib.templates import load_template +from ptodsl.tilelib.templates.a5._elementwise import ( + emit_scalar_binary_1d, + emit_scalar_binary_2d, + emit_scalar_fill_1d, + emit_scalar_fill_2d, + emit_unary_1d, + emit_unary_2d, + register_binary, +) + + +ELEMENTWISE_SCOPE_BY_FAMILY = { + "unary": frozenset( + { + "pto.tabs", + "pto.texp", + "pto.tlog", + "pto.tneg", + "pto.tnot", + "pto.trecip", + "pto.trelu", + "pto.trsqrt", + "pto.tsqrt", + } + ), + "tile_tile": frozenset( + { + "pto.tadd", + "pto.tand", + "pto.tdiv", + "pto.tfmod", + "pto.tmax", + "pto.tmin", + "pto.tmul", + "pto.tor", + "pto.tprelu", + "pto.trem", + "pto.tshl", + "pto.tshr", + "pto.tsub", + "pto.txor", + } + ), + "tile_scalar": frozenset( + { + "pto.tadds", + "pto.tands", + "pto.tdivs", + "pto.tfmods", + "pto.tlrelu", + "pto.tmaxs", + "pto.tmins", + "pto.tmuls", + "pto.tors", + "pto.trems", + "pto.tshls", + "pto.tshrs", + "pto.tsubs", + "pto.txors", + } + ), + "compare": frozenset({"pto.tcmp", "pto.tcmps"}), + "select": frozenset({"pto.tsel", "pto.tsels"}), + "conversion": frozenset({"pto.tcvt"}), + "scalar_fill": frozenset({"pto.texpands"}), +} + +ELEMENTWISE_SCOPE = frozenset().union( + *ELEMENTWISE_SCOPE_BY_FAMILY.values() +) + +# Map an operation to its reviewed reason if no flattened form is legal. +# Every scoped operation currently has at least one legal 1D candidate. +ELEMENTWISE_1D_EXCEPTIONS = {} # op -> (expected template name, expected vector op in the rendered MLIR) ELEMENTWISE = { @@ -21,6 +108,161 @@ "pto.tdiv": ("template_tdiv", "pto.vdiv"), } +PRODUCTION_UNARY_1D = { + "pto.tabs": ("template_tabs_1d", "template_tabs", "pto.vabs", "f32"), + "pto.texp": ("template_texp_1d", "template_texp", "pto.vexp", "f32"), + "pto.tneg": ("template_tneg_1d", "template_tneg", "pto.vneg", "i8"), + "pto.tnot": ("template_tnot_1d", "template_tnot", "pto.vnot", "ui16"), + "pto.trelu": ("template_trelu_1d", "template_trelu", "pto.vrelu", "i32"), + "pto.trsqrt": ( + "template_trsqrt_1d", + "template_trsqrt", + "pto.vsqrt", + "f16", + ), + "pto.tsqrt": ( + "template_tsqrt_1d", + "template_tsqrt", + "pto.vsqrt", + "f32", + ), +} + +PRODUCTION_BINARY_1D = { + "pto.tadd": ("template_tadd_1d", "template_tadd", "pto.vadd", "f32"), + "pto.tand": ("template_tand_1d", "template_tand", "pto.vand", "i8"), + "pto.tmax": ("template_tmax_1d", "template_tmax", "pto.vmax", "f16"), + "pto.tmin": ("template_tmin_1d", "template_tmin", "pto.vmin", "i16"), + "pto.tmul": ("template_tmul_1d", "template_tmul", "pto.vmul", "bf16"), + "pto.tor": ("template_tor_1d", "template_tor", "pto.vor", "ui8"), + "pto.tshl": ("template_tshl_1d", "template_tshl", "pto.vshl", "ui16"), + "pto.tshr": ("template_tshr_1d", "template_tshr", "pto.vshr", "ui32"), + "pto.tsub": ("template_tsub_1d", "template_tsub", "pto.vsub", "i32"), +} + +PRODUCTION_TEMP_BINARY_1D = { + "pto.tprelu": ( + "template_tprelu_1d", + "template_tprelu", + "pto.vprelu", + "f32", + ), + "pto.trem": ( + "template_trem_1d", + "template_trem", + "pto.vtrc", + "f32", + ), + "pto.txor": ( + "template_txor_1d", + "template_txor", + "pto.vxor", + "i16", + ), +} + +PRODUCTION_SCALAR_1D = { + "pto.tadds": ( + "template_tadds_1d", + "template_tadds", + "pto.vadds", + "f32", + "f32", + ), + "pto.tands": ( + "template_tands_1d", + "template_tands", + "pto.vand", + "i8", + "i8", + ), + "pto.tmaxs": ( + "template_tmaxs_1d", + "template_tmaxs", + "pto.vmaxs", + "f16", + "f16", + ), + "pto.tmins": ( + "template_tmins_1d", + "template_tmins", + "pto.vmins", + "ui16", + "ui16", + ), + "pto.tmuls": ( + "template_tmuls_1d", + "template_tmuls", + "pto.vmuls", + "f32", + "f32", + ), + "pto.tors": ( + "template_tors_1d", + "template_tors", + "pto.vor", + "ui8", + "ui8", + ), + "pto.tshls": ( + "template_tshls_1d", + "template_tshls", + "pto.vshls", + "ui16", + "i16", + ), + "pto.tshrs": ( + "template_tshrs_1d", + "template_tshrs", + "pto.vshrs", + "ui32", + "i16", + ), + "pto.tsubs": ( + "template_tsubs_1d", + "template_tsubs", + "pto.vsub", + "i32", + "i32", + ), +} + +PRODUCTION_SPECIAL_SCALAR_1D = { + "pto.tfmods": ( + "template_tfmods_1d", + "template_tfmods", + "pto.vtrc", + "f32", + "f32", + "scalar", + ), + "pto.tlrelu": ( + "template_tlrelu_1d", + "template_tlrelu", + "pto.vlrelu", + "f16", + "f32", + "slope", + ), +} + +PRODUCTION_TEMP_SCALAR_1D = { + "pto.trems": ( + "template_trems_1d", + "template_trems", + "pto.vtrc", + "f32", + ), + "pto.txors": ( + "template_txors_1d", + "template_txors", + "pto.vxor", + "i16", + ), +} + +PRODUCTION_FILL_DTYPES = ("i8", "i16", "i32", "f16", "bf16", "f32") + # Structured abstraction every elementwise template must preserve. SHARED_OPS = ["pto.tile_buf_addr", "!pto.ptr", "scf.for", "iter_args", "pto.plt_b32", "pto.vlds", "pto.vsts", "pto.tilelang.instance"] @@ -31,7 +273,432 @@ def _f32_specs(): return {"src0": spec, "src1": spec, "dst": spec} +def _test_template(*, name, loop_depth): + return tilelib.tile_template( + op=f"test.{name}", + target="a5", + name=name, + loop_depth=loop_depth, + register=False, + ) + + +@_test_template(name="test_unary_1d", loop_depth=1) +def _test_unary_1d(src: pto.Tile, dst: pto.Tile): + emit_unary_1d(src, dst, pto.vabs) + + +@_test_template(name="test_unary_2d", loop_depth=2) +def _test_unary_2d(src: pto.Tile, dst: pto.Tile): + emit_unary_2d(src, dst, pto.vabs) + + +_test_binary_1d = register_binary( + op="test.elementwise.binary", + name="test_binary_1d", + vector_op=pto.vadd, + dtypes=[("f32", "f32", "f32")], + traversal="1d", + priority=10, + candidate_id=99, +) + + +_test_binary_2d = register_binary( + op="test.elementwise.binary", + name="test_binary_2d", + vector_op=pto.vadd, + dtypes=[("f32", "f32", "f32")], + traversal="2d", + priority=0, + candidate_id=0, +) + + +@_test_template(name="test_scalar_1d", loop_depth=1) +def _test_scalar_1d(src: pto.Tile, scalar, dst: pto.Tile): + emit_scalar_binary_1d(src, scalar, dst, pto.vadds) + + +@_test_template(name="test_scalar_2d", loop_depth=2) +def _test_scalar_2d(src: pto.Tile, scalar, dst: pto.Tile): + emit_scalar_binary_2d(src, scalar, dst, pto.vadds) + + +@_test_template(name="test_fill_1d", loop_depth=1) +def _test_fill_1d(scalar, dst: pto.Tile): + emit_scalar_fill_1d(scalar, dst) + + +@_test_template(name="test_fill_2d", loop_depth=2) +def _test_fill_2d(scalar, dst: pto.Tile): + emit_scalar_fill_2d(scalar, dst) + + +_FAMILY_FORMS = ( + (_test_unary_1d, 1, "pto.vabs", ("src", "dst")), + (_test_unary_2d, 2, "pto.vabs", ("src", "dst")), + (_test_binary_1d, 1, "pto.vadd", ("src0", "src1", "dst")), + (_test_binary_2d, 2, "pto.vadd", ("src0", "src1", "dst")), + (_test_scalar_1d, 1, "pto.vadds", ("src", "scalar", "dst")), + (_test_scalar_2d, 2, "pto.vadds", ("src", "scalar", "dst")), + (_test_fill_1d, 1, "pto.vdup", ("scalar", "dst")), + (_test_fill_2d, 2, "pto.vdup", ("scalar", "dst")), +) + + +def _foundation_specs(param_names, *, shape=(4, 65), valid_shape=None): + tile = TileSpec( + shape=shape, + valid_shape=valid_shape, + dtype=ScalarType("f32"), + ) + scalar = ScalarSpec(dtype=ScalarType("f32")) + return { + name: scalar if name == "scalar" else tile + for name in param_names + } + + +def _unary_specs( + dtype_name, + *, + shape=(4, 65), + valid_shape=None, + compact_mode="null", +): + tile = TileSpec( + shape=shape, + valid_shape=valid_shape, + dtype=ScalarType(dtype_name), + compact_mode=compact_mode, + ) + return {"src": tile, "dst": tile} + + +def _binary_specs( + dtype_name, + *, + shape=(4, 65), + valid_shape=None, + compact_mode="null", +): + tile = TileSpec( + shape=shape, + valid_shape=valid_shape, + dtype=ScalarType(dtype_name), + compact_mode=compact_mode, + ) + return {"src0": tile, "src1": tile, "dst": tile} + + +def _binary_tmp_specs( + data_dtype, + *, + tmp_dtype=None, + shape=(4, 65), + valid_shape=None, + compact_mode="null", + tmp_compact_mode=None, +): + data = TileSpec( + shape=shape, + valid_shape=valid_shape, + dtype=ScalarType(data_dtype), + compact_mode=compact_mode, + ) + tmp = TileSpec( + shape=shape, + valid_shape=valid_shape, + dtype=ScalarType(tmp_dtype or data_dtype), + compact_mode=( + compact_mode + if tmp_compact_mode is None + else tmp_compact_mode + ), + ) + return {"src0": data, "src1": data, "tmp": tmp, "dst": data} + + +def _scalar_specs( + data_dtype, + *, + scalar_dtype=None, + shape=(4, 65), + valid_shape=None, + compact_mode="null", +): + tile = TileSpec( + shape=shape, + valid_shape=valid_shape, + dtype=ScalarType(data_dtype), + compact_mode=compact_mode, + ) + scalar = ScalarSpec( + dtype=ScalarType(scalar_dtype or data_dtype), + value=1, + ) + return {"src": tile, "scalar": scalar, "dst": tile} + + +def _named_scalar_specs( + data_dtype, + *, + scalar_dtype=None, + scalar_name="scalar", + shape=(4, 65), + valid_shape=None, + compact_mode="null", +): + specs = _scalar_specs( + data_dtype, + scalar_dtype=scalar_dtype, + shape=shape, + valid_shape=valid_shape, + compact_mode=compact_mode, + ) + specs[scalar_name] = specs.pop("scalar") + return specs + + +def _scalar_tmp_specs( + data_dtype, + *, + scalar_dtype=None, + shape=(4, 65), + valid_shape=None, + compact_mode="null", + tmp_compact_mode=None, +): + specs = _scalar_specs( + data_dtype, + scalar_dtype=scalar_dtype, + shape=shape, + valid_shape=valid_shape, + compact_mode=compact_mode, + ) + specs["tmp"] = TileSpec( + shape=shape, + valid_shape=valid_shape, + dtype=ScalarType(data_dtype), + compact_mode=( + compact_mode + if tmp_compact_mode is None + else tmp_compact_mode + ), + ) + return { + "src": specs["src"], + "scalar": specs["scalar"], + "tmp": specs["tmp"], + "dst": specs["dst"], + } + + +def _fill_specs( + dtype_name, + *, + shape=(4, 65), + valid_shape=None, + compact_mode="null", +): + return { + "scalar": ScalarSpec( + dtype=ScalarType(dtype_name), + value=1, + ), + "dst": TileSpec( + shape=shape, + valid_shape=valid_shape, + dtype=ScalarType(dtype_name), + compact_mode=compact_mode, + ), + } + + +def _compare_specs( + op, + dtype_name, + *, + src_shape=(1, 256), + src_valid_shape=None, + dst_shape=None, + dst_valid_shape=None, + src_compact_mode="null", + dst_compact_mode="null", +): + src = TileSpec( + shape=src_shape, + valid_shape=src_valid_shape, + dtype=ScalarType(dtype_name), + memory_space="vec", + compact_mode=src_compact_mode, + ) + predicate_dtype = "i8" if op == "pto.tcmp" else "ui8" + if dst_shape is None: + dst_shape = src_shape if op == "pto.tcmp" else (src_shape[0], 32) + dst = TileSpec( + shape=dst_shape, + valid_shape=dst_valid_shape, + dtype=ScalarType(predicate_dtype), + memory_space="vec", + compact_mode=dst_compact_mode, + ) + if op == "pto.tcmp": + return {"src0": src, "src1": src, "dst": dst} + return { + "src": src, + "scalar": ScalarSpec(dtype=ScalarType(dtype_name), value=1), + "dst": dst, + } + + +def _select_specs( + op, + data_dtype, + *, + mask_dtype="i8", + data_shape=(1, 256), + data_valid_shape=None, + mask_shape=None, + mask_valid_shape=None, + tmp_shape=(1, 64), + data_compact_mode="null", + mask_compact_mode="null", + tmp_compact_mode="null", +): + bytewidth = {"i8": 1, "i16": 2, "i32": 4}[mask_dtype] + if mask_shape is None: + mask_shape = (data_shape[0], 32 // bytewidth) + mask = TileSpec( + shape=mask_shape, + valid_shape=mask_valid_shape, + dtype=ScalarType(mask_dtype), + compact_mode=mask_compact_mode, + ) + data = TileSpec( + shape=data_shape, + valid_shape=data_valid_shape, + dtype=ScalarType(data_dtype), + compact_mode=data_compact_mode, + ) + tmp = TileSpec( + shape=tmp_shape, + dtype=ScalarType(data_dtype), + compact_mode=tmp_compact_mode, + ) + if op == "pto.tsel": + return { + "mask": mask, + "src0": data, + "src1": data, + "tmp": tmp, + "dst": data, + } + return { + "mask": mask, + "src": data, + "tmp": tmp, + "scalar": ScalarSpec(dtype=ScalarType(data_dtype), value=1), + "dst": data, + } + + class TileLibElementwiseTest(unittest.TestCase): + def test_issue_scope_catalog_has_ranked_1d_and_2d_coverage(self): + family_sizes = sum( + len(operations) + for operations in ELEMENTWISE_SCOPE_BY_FAMILY.values() + ) + self.assertEqual(family_sizes, 43) + self.assertEqual(len(ELEMENTWISE_SCOPE), family_sizes) + self.assertLessEqual( + set(ELEMENTWISE_1D_EXCEPTIONS), + ELEMENTWISE_SCOPE, + ) + + for op in sorted(ELEMENTWISE_SCOPE): + with self.subTest(op=op): + self.assertTrue(load_template(op, "a5")) + candidates = tilelib.default_registry().lookup(op, "a5") + self.assertTrue(candidates) + + candidate_ids = [ + candidate.metadata.id for candidate in candidates + ] + self.assertEqual( + len(candidate_ids), + len(set(candidate_ids)), + "candidate IDs must remain unique within an operation", + ) + + by_loop_depth = { + loop_depth: [ + candidate + for candidate in candidates + if candidate.metadata.loop_depth == loop_depth + ] + for loop_depth in (1, 2) + } + self.assertTrue( + by_loop_depth[2], + "every scoped operation needs a general 2D fallback", + ) + + exception = ELEMENTWISE_1D_EXCEPTIONS.get(op) + if exception is not None: + self.assertTrue(exception.strip()) + self.assertFalse(by_loop_depth[1]) + continue + + self.assertTrue( + by_loop_depth[1], + "operation needs a 1D candidate or a documented exception", + ) + self.assertGreater( + min( + candidate.metadata.priority + for candidate in by_loop_depth[1] + ), + max( + candidate.metadata.priority + for candidate in by_loop_depth[2] + ), + "every legal 1D form must rank ahead of every 2D fallback", + ) + + def test_shared_traversals_use_native_python_for_syntax(self): + tree = ast.parse( + Path(elementwise.__file__).read_text(encoding="utf-8") + ) + functions = { + node.name: node + for node in tree.body + if isinstance(node, ast.FunctionDef) + } + + for function_name, expected_loops in ( + ("emit_elementwise_1d", 1), + ("emit_elementwise_2d", 2), + ): + with self.subTest(function=function_name): + function = functions[function_name] + self.assertEqual( + sum( + isinstance(node, ast.For) + for node in ast.walk(function) + ), + expected_loops, + ) + self.assertFalse( + any( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "for_" + for node in ast.walk(function) + ) + ) + def test_each_op_selects_and_renders(self): for op, (name, vop) in ELEMENTWISE.items(): with self.subTest(op=op): @@ -44,6 +711,1221 @@ def test_each_op_selects_and_renders(self): self.assertIn(shared, mlir) self.assertNotIn("pto.castptr", mlir) # structured, not bare-pointer + def test_shared_family_forms_expose_expected_loop_depth(self): + for descriptor, expected_depth, vector_op, params in _FAMILY_FORMS: + with self.subTest(template=descriptor.name): + self.assertEqual(descriptor.metadata.loop_depth, expected_depth) + mlir = descriptor.specialize( + **_foundation_specs(params) + ).mlir_text() + self.assertEqual(mlir.count("scf.for"), expected_depth) + self.assertIn(vector_op, mlir) + self.assertIn("iter_args", mlir) + self.assertIn("pto.plt_b32", mlir) + + def test_flattened_form_uses_one_total_element_loop_for_aligned_and_tail_shapes(self): + for shape in ((4, 64), (4, 65), (1, 65)): + with self.subTest(shape=shape): + mlir = _test_binary_1d.specialize( + **_foundation_specs( + ("src0", "src1", "dst"), + shape=shape, + ) + ).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn("arith.muli", mlir) + self.assertIn("iter_args", mlir) + self.assertNotIn("memref.subview", mlir) + + def test_rowwise_form_retains_row_and_column_loops(self): + mlir = _test_binary_2d.specialize( + **_foundation_specs( + ("src0", "src1", "dst"), + shape=(4, 65), + valid_shape=(4, 63), + ) + ).mlir_text() + # Tile slicing may lower either to memref.subview or to explicit + # row-stride pointer arithmetic. Both retain the same 2D traversal. + self.assertEqual(mlir.count("scf.for"), 2) + self.assertIn("iter_args", mlir) + + def test_registration_helper_selects_1d_only_for_legal_shapes(self): + contiguous = _foundation_specs(("src0", "src1", "dst"), shape=(4, 65)) + partial = _foundation_specs( + ("src0", "src1", "dst"), + shape=(4, 65), + valid_shape=(4, 63), + ) + + selected_1d = select("test.elementwise.binary", "a5", contiguous) + selected_2d = select("test.elementwise.binary", "a5", partial) + + self.assertEqual(selected_1d.name, "test_binary_1d") + self.assertEqual(selected_1d.metadata.loop_depth, 1) + self.assertEqual(selected_2d.name, "test_binary_2d") + self.assertEqual(selected_2d.metadata.loop_depth, 2) + + def test_production_unary_ops_register_preferred_1d_and_fallback_2d(self): + for op, ( + name_1d, + name_2d, + vector_op, + dtype_name, + ) in PRODUCTION_UNARY_1D.items(): + with self.subTest(op=op): + specs = _unary_specs(dtype_name) + candidates = legal_candidates(op, "a5", specs) + + self.assertEqual( + [candidate.name for candidate in candidates], + [name_1d, name_2d], + ) + self.assertEqual( + [candidate.metadata.id for candidate in candidates], + [1, 0], + ) + self.assertEqual( + [candidate.metadata.loop_depth for candidate in candidates], + [1, 2], + ) + + mlir = candidates[0].specialize(**specs).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn(vector_op, mlir) + self.assertNotIn("memref.subview", mlir) + + def test_production_unary_selection_uses_shared_legality_rule(self): + shapes = ( + ("contiguous multi-row", (4, 65), None, "null", True), + ("contiguous single-row", (4, 65), (1, 63), "null", True), + ("partial multi-row", (4, 65), (4, 63), "null", False), + ("stride gap", (4, 65), None, 2, False), + ) + for op, ( + name_1d, + name_2d, + _, + dtype_name, + ) in PRODUCTION_UNARY_1D.items(): + for label, shape, valid_shape, compact_mode, expect_1d in shapes: + with self.subTest(op=op, case=label): + specs = _unary_specs( + dtype_name, + shape=shape, + valid_shape=valid_shape, + compact_mode=compact_mode, + ) + selected = select(op, "a5", specs) + self.assertEqual( + selected.name, + name_1d if expect_1d else name_2d, + ) + self.assertEqual( + selected.metadata.loop_depth, + 1 if expect_1d else 2, + ) + + def test_tlog_precision_modes_register_distinct_1d_and_2d_candidates(self): + specs = _unary_specs("f32") + cases = ( + ( + "default", + ["template_tlog_1d", "template_tlog"], + [2, 0], + ("pto.vln",), + ), + ( + "high_precision", + [ + "template_tlog_high_precision_1d", + "template_tlog_high_precision", + ], + [3, 1], + ( + "pto.vcmps", + "pto.vmuls", + "pto.vsel", + "pto.vln", + "pto.vadds", + ), + ), + ) + for precision_type, expected_names, expected_ids, vector_ops in cases: + with self.subTest(precision_type=precision_type): + context_attrs = {"precisionType": precision_type} + candidates = legal_candidates( + "pto.tlog", + "a5", + specs, + context_attrs=context_attrs, + ) + + self.assertEqual( + [candidate.name for candidate in candidates], + expected_names, + ) + self.assertEqual( + [candidate.metadata.id for candidate in candidates], + expected_ids, + ) + self.assertEqual( + [candidate.metadata.loop_depth for candidate in candidates], + [1, 2], + ) + + mlir = candidates[0].specialize( + context_attrs=context_attrs, + **specs, + ).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertNotIn("memref.subview", mlir) + for vector_op in vector_ops: + self.assertIn(vector_op, mlir) + + def test_trecip_registers_local_computation_with_shared_traversals(self): + specs = _unary_specs("f16") + candidates = legal_candidates("pto.trecip", "a5", specs) + + self.assertEqual( + [candidate.name for candidate in candidates], + ["template_trecip_1d", "template_trecip"], + ) + self.assertEqual( + [candidate.metadata.id for candidate in candidates], + [1, 0], + ) + self.assertEqual( + [candidate.metadata.loop_depth for candidate in candidates], + [1, 2], + ) + + mlir = candidates[0].specialize(**specs).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertNotIn("memref.subview", mlir) + self.assertIn("pto.vbr", mlir) + self.assertIn("pto.vdiv", mlir) + + def test_specialized_unary_selection_uses_shared_legality_rule(self): + shapes = ( + ("contiguous multi-row", (4, 65), None, "null", True), + ("contiguous single-row", (4, 65), (1, 63), "null", True), + ("partial multi-row", (4, 65), (4, 63), "null", False), + ("stride gap", (4, 65), None, 2, False), + ) + operations = ( + ( + "pto.tlog", + "template_tlog_high_precision_1d", + "template_tlog_high_precision", + {"precisionType": "high_precision"}, + ), + ( + "pto.trecip", + "template_trecip_1d", + "template_trecip", + None, + ), + ) + for op, name_1d, name_2d, context_attrs in operations: + for label, shape, valid_shape, compact_mode, expect_1d in shapes: + with self.subTest(op=op, case=label): + specs = _unary_specs( + "f32", + shape=shape, + valid_shape=valid_shape, + compact_mode=compact_mode, + ) + selected = select( + op, + "a5", + specs, + context_attrs=context_attrs, + ) + self.assertEqual( + selected.name, + name_1d if expect_1d else name_2d, + ) + self.assertEqual( + selected.metadata.loop_depth, + 1 if expect_1d else 2, + ) + + def test_production_binary_ops_register_preferred_1d_and_fallback_2d(self): + for op, ( + name_1d, + name_2d, + vector_op, + dtype_name, + ) in PRODUCTION_BINARY_1D.items(): + with self.subTest(op=op): + specs = _binary_specs(dtype_name) + candidates = legal_candidates(op, "a5", specs) + + self.assertEqual( + [candidate.name for candidate in candidates], + [name_1d, name_2d], + ) + self.assertEqual( + [candidate.metadata.id for candidate in candidates], + [1, 0], + ) + self.assertEqual( + [candidate.metadata.loop_depth for candidate in candidates], + [1, 2], + ) + self.assertEqual( + candidates[0].metadata.dtypes, + candidates[1].metadata.dtypes, + ) + + mlir = candidates[0].specialize(**specs).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn(vector_op, mlir) + self.assertNotIn("memref.subview", mlir) + + def test_production_binary_selection_uses_shared_legality_rule(self): + shapes = ( + ("contiguous multi-row", (4, 65), None, "null", True), + ("contiguous single-row", (4, 65), (1, 63), "null", True), + ("partial multi-row", (4, 65), (4, 63), "null", False), + ("stride gap", (4, 65), None, 2, False), + ) + for op, ( + name_1d, + name_2d, + _, + dtype_name, + ) in PRODUCTION_BINARY_1D.items(): + for label, shape, valid_shape, compact_mode, expect_1d in shapes: + with self.subTest(op=op, case=label): + specs = _binary_specs( + dtype_name, + shape=shape, + valid_shape=valid_shape, + compact_mode=compact_mode, + ) + selected = select(op, "a5", specs) + self.assertEqual( + selected.name, + name_1d if expect_1d else name_2d, + ) + self.assertEqual( + selected.metadata.loop_depth, + 1 if expect_1d else 2, + ) + + def test_tdiv_precision_modes_share_ranked_traversal_candidates(self): + specs = _binary_specs("f32") + cases = ( + ("default", ("pto.vdiv",)), + ( + "high_precision", + ( + "pto.vdiv", + "pto.vbitcast", + "pto.vcmp", + "pto.vsel", + ), + ), + ) + for precision_type, vector_ops in cases: + with self.subTest(precision_type=precision_type): + context_attrs = {"precisionType": precision_type} + candidates = legal_candidates( + "pto.tdiv", + "a5", + specs, + context_attrs=context_attrs, + ) + + self.assertEqual( + [candidate.name for candidate in candidates], + ["template_tdiv_1d", "template_tdiv"], + ) + self.assertEqual( + [candidate.metadata.id for candidate in candidates], + [1, 0], + ) + self.assertEqual( + [candidate.metadata.loop_depth for candidate in candidates], + [1, 2], + ) + + mlir = candidates[0].specialize( + context_attrs=context_attrs, + **specs, + ).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertNotIn("memref.subview", mlir) + for vector_op in vector_ops: + self.assertIn(vector_op, mlir) + + def test_tfmod_1d_preserves_dtype_specific_remainder_computation(self): + cases = ( + ("f32", True), + ("f16", True), + ("i16", False), + ("ui16", False), + ) + for dtype_name, expects_truncation in cases: + with self.subTest(dtype=dtype_name): + specs = _binary_specs(dtype_name) + candidates = legal_candidates("pto.tfmod", "a5", specs) + + self.assertEqual( + [candidate.name for candidate in candidates], + ["template_tfmod_1d", "template_tfmod"], + ) + self.assertEqual( + [candidate.metadata.id for candidate in candidates], + [1, 0], + ) + self.assertEqual( + [candidate.metadata.loop_depth for candidate in candidates], + [1, 2], + ) + + mlir = candidates[0].specialize(**specs).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertNotIn("memref.subview", mlir) + for vector_op in ("pto.vdiv", "pto.vmul", "pto.vsub"): + self.assertIn(vector_op, mlir) + if expects_truncation: + self.assertIn("pto.vtrc", mlir) + else: + self.assertNotIn("pto.vtrc", mlir) + + def test_specialized_binary_selection_uses_shared_legality_rule(self): + shapes = ( + ("contiguous multi-row", (4, 65), None, "null", True), + ("contiguous single-row", (4, 65), (1, 63), "null", True), + ("partial multi-row", (4, 65), (4, 63), "null", False), + ("stride gap", (4, 65), None, 2, False), + ) + operations = ( + ("pto.tdiv", "template_tdiv_1d", "template_tdiv"), + ("pto.tfmod", "template_tfmod_1d", "template_tfmod"), + ) + for op, name_1d, name_2d in operations: + for label, shape, valid_shape, compact_mode, expect_1d in shapes: + with self.subTest(op=op, case=label): + specs = _binary_specs( + "f32", + shape=shape, + valid_shape=valid_shape, + compact_mode=compact_mode, + ) + selected = select(op, "a5", specs) + self.assertEqual( + selected.name, + name_1d if expect_1d else name_2d, + ) + self.assertEqual( + selected.metadata.loop_depth, + 1 if expect_1d else 2, + ) + + def test_temporary_binary_ops_register_preferred_1d_and_fallback_2d(self): + for op, ( + name_1d, + name_2d, + vector_op, + dtype_name, + ) in PRODUCTION_TEMP_BINARY_1D.items(): + with self.subTest(op=op): + specs = _binary_tmp_specs(dtype_name) + candidates = legal_candidates(op, "a5", specs) + + self.assertEqual( + [candidate.name for candidate in candidates], + [name_1d, name_2d], + ) + self.assertEqual( + [candidate.metadata.id for candidate in candidates], + [1, 0], + ) + self.assertEqual( + [candidate.metadata.loop_depth for candidate in candidates], + [1, 2], + ) + self.assertEqual( + candidates[0].metadata.dtypes, + candidates[1].metadata.dtypes, + ) + + mlir = candidates[0].specialize(**specs).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn(vector_op, mlir) + self.assertNotIn("memref.subview", mlir) + + def test_temporary_binary_selection_uses_shared_legality_rule(self): + shapes = ( + ("contiguous multi-row", (4, 65), None, "null", True), + ("contiguous single-row", (4, 65), (1, 63), "null", True), + ("partial multi-row", (4, 65), (4, 63), "null", False), + ("stride gap", (4, 65), None, 2, False), + ) + for op, ( + name_1d, + name_2d, + _, + dtype_name, + ) in PRODUCTION_TEMP_BINARY_1D.items(): + for label, shape, valid_shape, compact_mode, expect_1d in shapes: + with self.subTest(op=op, case=label): + specs = _binary_tmp_specs( + dtype_name, + shape=shape, + valid_shape=valid_shape, + compact_mode=compact_mode, + ) + selected = select(op, "a5", specs) + self.assertEqual( + selected.name, + name_1d if expect_1d else name_2d, + ) + self.assertEqual( + selected.metadata.loop_depth, + 1 if expect_1d else 2, + ) + + def test_temporary_tile_alone_can_disqualify_1d(self): + for op, ( + _, + name_2d, + _, + dtype_name, + ) in PRODUCTION_TEMP_BINARY_1D.items(): + with self.subTest(op=op): + specs = _binary_tmp_specs( + dtype_name, + tmp_compact_mode=2, + ) + candidates = legal_candidates(op, "a5", specs) + + self.assertEqual( + [candidate.name for candidate in candidates], + [name_2d], + ) + self.assertEqual(candidates[0].metadata.loop_depth, 2) + + def test_tprelu_1d_accepts_i8_temporary_representation(self): + specs = _binary_tmp_specs("f32", tmp_dtype="i8") + candidates = legal_candidates("pto.tprelu", "a5", specs) + + self.assertEqual( + [candidate.name for candidate in candidates], + ["template_tprelu_1d", "template_tprelu"], + ) + mlir = candidates[0].specialize(**specs).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn("pto.vprelu", mlir) + + def test_trem_1d_preserves_floor_remainder_computation(self): + specs = _binary_tmp_specs("f32") + descriptor = select("pto.trem", "a5", specs) + mlir = descriptor.specialize(**specs).mlir_text() + + self.assertEqual(descriptor.name, "template_trem_1d") + self.assertEqual(mlir.count("scf.for"), 1) + for vector_op in ( + "pto.vdiv", + "pto.vtrc", + "pto.vmul", + "pto.vsub", + ): + self.assertIn(vector_op, mlir) + + def test_production_scalar_ops_register_preferred_1d_and_fallback_2d(self): + for op, ( + name_1d, + name_2d, + vector_op, + data_dtype, + scalar_dtype, + ) in PRODUCTION_SCALAR_1D.items(): + with self.subTest(op=op): + specs = _scalar_specs( + data_dtype, + scalar_dtype=scalar_dtype, + ) + candidates = legal_candidates(op, "a5", specs) + + self.assertEqual( + [candidate.name for candidate in candidates], + [name_1d, name_2d], + ) + self.assertEqual( + [candidate.metadata.id for candidate in candidates], + [1, 0], + ) + self.assertEqual( + [candidate.metadata.loop_depth for candidate in candidates], + [1, 2], + ) + self.assertEqual( + candidates[0].metadata.dtypes, + candidates[1].metadata.dtypes, + ) + + for shape in ((4, 64), (4, 65)): + with self.subTest(op=op, shape=shape): + render_specs = _scalar_specs( + data_dtype, + scalar_dtype=scalar_dtype, + shape=shape, + ) + mlir = candidates[0].specialize( + **render_specs + ).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn(vector_op, mlir) + self.assertNotIn("memref.subview", mlir) + + def test_production_scalar_selection_uses_shared_legality_rule(self): + shapes = ( + ("contiguous multi-row", (4, 65), None, "null", True), + ("contiguous single-row", (4, 65), (1, 63), "null", True), + ("partial multi-row", (4, 65), (4, 63), "null", False), + ("stride gap", (4, 65), None, 2, False), + ) + for op, ( + name_1d, + name_2d, + _, + data_dtype, + scalar_dtype, + ) in PRODUCTION_SCALAR_1D.items(): + for label, shape, valid_shape, compact_mode, expect_1d in shapes: + with self.subTest(op=op, case=label): + specs = _scalar_specs( + data_dtype, + scalar_dtype=scalar_dtype, + shape=shape, + valid_shape=valid_shape, + compact_mode=compact_mode, + ) + selected = select(op, "a5", specs) + self.assertEqual( + selected.name, + name_1d if expect_1d else name_2d, + ) + self.assertEqual( + selected.metadata.loop_depth, + 1 if expect_1d else 2, + ) + + def test_tsubs_shared_traversal_preserves_broadcast_subtraction(self): + specs = _scalar_specs("f32") + candidates = legal_candidates("pto.tsubs", "a5", specs) + + for candidate, expected_loops in zip(candidates, (1, 2)): + with self.subTest(candidate=candidate.name): + mlir = candidate.specialize(**specs).mlir_text() + self.assertEqual(mlir.count("scf.for"), expected_loops) + self.assertIn("pto.vbr", mlir) + self.assertIn("pto.vsub", mlir) + + def test_tdivs_operand_forms_preserve_candidate_ids(self): + cases = ( + ( + _scalar_specs("f32"), + ( + ("template_tdivs_tile_scalar_1d", 2, 1), + ("template_tdivs_tile_scalar", 0, 2), + ), + ), + ( + { + "scalar": ScalarSpec( + dtype=ScalarType("f32"), + value=1, + ), + "src": TileSpec( + shape=(4, 65), + dtype=ScalarType("f32"), + ), + "dst": TileSpec( + shape=(4, 65), + dtype=ScalarType("f32"), + ), + }, + ( + ("template_tdivs_scalar_tile_1d", 3, 1), + ("template_tdivs_scalar_tile", 1, 2), + ), + ), + ) + for specs, candidates in cases: + with self.subTest(operand_order=tuple(specs)): + for name, candidate_id, loop_depth in candidates: + descriptor = select( + "pto.tdivs", + "a5", + specs, + candidate_id=name, + ) + self.assertEqual(descriptor.metadata.id, candidate_id) + self.assertEqual( + descriptor.metadata.loop_depth, + loop_depth, + ) + + def test_tdivs_1d_preserves_precision_and_operand_order(self): + cases = ( + ( + _scalar_specs("f32"), + "template_tdivs_tile_scalar_1d", + ), + ( + { + "scalar": ScalarSpec( + dtype=ScalarType("f32"), + value=1, + ), + "src": TileSpec( + shape=(4, 65), + dtype=ScalarType("f32"), + ), + "dst": TileSpec( + shape=(4, 65), + dtype=ScalarType("f32"), + ), + }, + "template_tdivs_scalar_tile_1d", + ), + ) + for specs, expected_name in cases: + for precision_type in ("default", "high_precision"): + with self.subTest( + operand_order=tuple(specs), + precision_type=precision_type, + ): + context_attrs = {"precisionType": precision_type} + descriptor = select( + "pto.tdivs", + "a5", + specs, + context_attrs=context_attrs, + candidate_id=expected_name, + ) + mlir = descriptor.specialize( + context_attrs=context_attrs, + **specs, + ).mlir_text() + + self.assertEqual(descriptor.name, expected_name) + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn("pto.vbr", mlir) + self.assertIn("pto.vdiv", mlir) + if precision_type == "high_precision": + self.assertIn("pto.vbitcast", mlir) + self.assertIn("pto.vcmp", mlir) + + def test_tdivs_ineligible_ranges_reject_1d_candidates(self): + specs = _scalar_specs( + "f32", + shape=(4, 65), + valid_shape=(4, 63), + ) + candidates = legal_candidates("pto.tdivs", "a5", specs) + + self.assertNotIn( + "template_tdivs_tile_scalar_1d", + [candidate.name for candidate in candidates], + ) + self.assertNotIn( + "template_tdivs_scalar_tile_1d", + [candidate.name for candidate in candidates], + ) + self.assertTrue( + all( + candidate.metadata.loop_depth == 2 + for candidate in candidates + ) + ) + + def test_specialized_scalar_ops_use_shared_legality_rule(self): + shapes = ( + ("contiguous multi-row", (4, 65), None, "null", True), + ("contiguous single-row", (4, 65), (1, 63), "null", True), + ("partial multi-row", (4, 65), (4, 63), "null", False), + ("stride gap", (4, 65), None, 2, False), + ) + for op, ( + name_1d, + name_2d, + vector_op, + data_dtype, + scalar_dtype, + scalar_name, + ) in PRODUCTION_SPECIAL_SCALAR_1D.items(): + for label, shape, valid_shape, compact_mode, expect_1d in shapes: + with self.subTest(op=op, case=label): + specs = _named_scalar_specs( + data_dtype, + scalar_dtype=scalar_dtype, + scalar_name=scalar_name, + shape=shape, + valid_shape=valid_shape, + compact_mode=compact_mode, + ) + candidates = legal_candidates(op, "a5", specs) + self.assertEqual( + candidates[0].name, + name_1d if expect_1d else name_2d, + ) + self.assertEqual( + candidates[0].metadata.loop_depth, + 1 if expect_1d else 2, + ) + + mlir = candidates[0].specialize(**specs).mlir_text() + self.assertIn(vector_op, mlir) + + def test_temporary_scalar_ops_include_tmp_in_1d_legality(self): + for op, ( + name_1d, + name_2d, + vector_op, + data_dtype, + ) in PRODUCTION_TEMP_SCALAR_1D.items(): + with self.subTest(op=op, case="contiguous"): + specs = _scalar_tmp_specs(data_dtype) + candidates = legal_candidates(op, "a5", specs) + self.assertEqual( + [candidate.name for candidate in candidates], + [name_1d, name_2d], + ) + self.assertEqual( + [candidate.metadata.id for candidate in candidates], + [1, 0], + ) + self.assertEqual( + [candidate.metadata.loop_depth for candidate in candidates], + [1, 2], + ) + mlir = candidates[0].specialize(**specs).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn(vector_op, mlir) + + with self.subTest(op=op, case="temporary stride gap"): + specs = _scalar_tmp_specs( + data_dtype, + tmp_compact_mode=2, + ) + candidates = legal_candidates(op, "a5", specs) + self.assertEqual( + [candidate.name for candidate in candidates], + [name_2d], + ) + self.assertEqual(candidates[0].metadata.loop_depth, 2) + + def test_scalar_remainder_1d_preserves_instruction_sequences(self): + for op, specs in ( + ("pto.tfmods", _scalar_specs("f32")), + ("pto.trems", _scalar_tmp_specs("f32")), + ): + with self.subTest(op=op): + descriptor = select(op, "a5", specs) + mlir = descriptor.specialize(**specs).mlir_text() + + self.assertEqual(descriptor.metadata.loop_depth, 1) + for vector_op in ( + "pto.vbr", + "pto.vdiv", + "pto.vtrc", + "pto.vmuls", + "pto.vsub", + ): + self.assertIn(vector_op, mlir) + + def test_texpands_registers_preferred_1d_and_fallback_2d(self): + for dtype_name in PRODUCTION_FILL_DTYPES: + with self.subTest(dtype=dtype_name): + specs = _fill_specs(dtype_name) + candidates = legal_candidates("pto.texpands", "a5", specs) + + self.assertEqual( + [candidate.name for candidate in candidates], + ["template_texpands_1d", "template_texpands"], + ) + self.assertEqual( + [candidate.metadata.id for candidate in candidates], + [1, 0], + ) + self.assertEqual( + [candidate.metadata.loop_depth for candidate in candidates], + [1, 2], + ) + self.assertEqual( + candidates[0].metadata.dtypes, + candidates[1].metadata.dtypes, + ) + + for shape in ((4, 64), (4, 65)): + with self.subTest(dtype=dtype_name, shape=shape): + mlir = candidates[0].specialize( + **_fill_specs(dtype_name, shape=shape) + ).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn("pto.vdup", mlir) + self.assertIn("pto.vsts", mlir) + self.assertNotIn("pto.vlds", mlir) + self.assertNotIn("memref.subview", mlir) + + def test_texpands_selection_uses_destination_legality(self): + cases = ( + ("contiguous multi-row", (4, 65), None, "null", True), + ("contiguous single-row", (4, 65), (1, 63), "null", True), + ("partial multi-row", (4, 65), (4, 63), "null", False), + ("stride gap", (4, 65), None, 2, False), + ) + for label, shape, valid_shape, compact_mode, expect_1d in cases: + with self.subTest(case=label): + specs = _fill_specs( + "f32", + shape=shape, + valid_shape=valid_shape, + compact_mode=compact_mode, + ) + selected = select("pto.texpands", "a5", specs) + + self.assertEqual( + selected.name, + "template_texpands_1d" + if expect_1d + else "template_texpands", + ) + self.assertEqual( + selected.metadata.loop_depth, + 1 if expect_1d else 2, + ) + + def test_compare_family_registers_preferred_1d_and_fallback_2d(self): + cases = ( + ("pto.tcmp", "template_tcmp_1d", "template_tcmp", "pto.vcmp"), + ( + "pto.tcmps", + "template_tcmps_1d", + "template_tcmps", + "pto.vcmps", + ), + ) + for op, name_1d, name_2d, vector_op in cases: + for dtype_name, expected_dist in ( + ("f32", "PK"), + ("f16", "PK"), + ("i8", "NORM"), + ): + with self.subTest(op=op, dtype=dtype_name): + specs = _compare_specs(op, dtype_name) + candidates = legal_candidates(op, "a5", specs) + self.assertEqual( + [candidate.name for candidate in candidates], + [name_1d, name_2d], + ) + self.assertEqual( + [candidate.metadata.id for candidate in candidates], + [1, 0], + ) + self.assertEqual( + [candidate.metadata.loop_depth for candidate in candidates], + [1, 2], + ) + + mlir = candidates[0].specialize( + context_attrs={"cmp_mode": "lt"}, + **specs, + ).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn(vector_op, mlir) + self.assertIn(expected_dist, mlir) + self.assertIn('"lt"', mlir) + + def test_compare_family_uses_predicate_specific_fallbacks(self): + cases = ( + ( + "tcmp multi-row predicate padding", + "pto.tcmp", + _compare_specs( + "pto.tcmp", + "f32", + src_shape=(4, 128), + ), + "template_tcmp", + ), + ( + "tcmps source row tail", + "pto.tcmps", + _compare_specs( + "pto.tcmps", + "i8", + src_shape=(4, 128), + dst_shape=(4, 32), + ), + "template_tcmps", + ), + ( + "tcmps predicate row padding", + "pto.tcmps", + _compare_specs( + "pto.tcmps", + "f32", + src_shape=(4, 256), + dst_shape=(4, 64), + ), + "template_tcmps", + ), + ( + "tcmps stride gap", + "pto.tcmps", + _compare_specs( + "pto.tcmps", + "f32", + src_shape=(4, 256), + dst_shape=(4, 32), + dst_compact_mode=2, + ), + "template_tcmps", + ), + ) + for label, op, specs, fallback_name in cases: + with self.subTest(case=label): + candidates = legal_candidates(op, "a5", specs) + self.assertEqual( + [candidate.name for candidate in candidates], + [fallback_name], + ) + self.assertEqual(candidates[0].metadata.loop_depth, 2) + + def test_tcmps_aligned_dense_multi_row_selects_flattened_body(self): + specs = _compare_specs( + "pto.tcmps", + "f32", + src_shape=(4, 256), + dst_shape=(4, 32), + ) + selected = select("pto.tcmps", "a5", specs) + mlir = selected.specialize(**specs).mlir_text() + + self.assertEqual(selected.name, "template_tcmps_1d") + self.assertEqual(selected.metadata.loop_depth, 1) + self.assertEqual(mlir.count("scf.for"), 1) + + def test_tcmps_f32_fallback_is_genuinely_row_wise(self): + specs = _compare_specs( + "pto.tcmps", + "f32", + src_shape=(4, 64), + dst_shape=(4, 32), + ) + selected = select("pto.tcmps", "a5", specs) + mlir = selected.specialize(**specs).mlir_text() + + self.assertEqual(selected.name, "template_tcmps") + self.assertEqual(selected.metadata.loop_depth, 2) + self.assertEqual(mlir.count("scf.for"), 2) + + def test_tcmps_fallback_keeps_flat_packed_predicate_rows(self): + specs = _compare_specs( + "pto.tcmps", + "f16", + src_shape=(4, 256), + dst_shape=(4, 64), + dst_valid_shape=(4, 32), + ) + selected = select("pto.tcmps", "a5", specs) + mlir = selected.specialize(**specs).mlir_text() + + self.assertEqual(selected.name, "template_tcmps") + self.assertEqual(selected.metadata.loop_depth, 2) + self.assertIn("arith.constant 16 : index", mlir) + self.assertNotIn("arith.constant 64 : index", mlir) + + def test_select_family_registers_preferred_1d_and_fallback_2d(self): + cases = ( + ("pto.tsel", "f32", "i8", "US"), + ("pto.tsel", "f16", "i8", "US"), + ("pto.tsel", "i8", "i8", "NORM"), + ("pto.tsels", "f32", "i32", "US"), + ("pto.tsels", "f16", "i16", "US"), + ("pto.tsels", "i8", "i8", "NORM"), + ) + for op, data_dtype, mask_dtype, expected_dist in cases: + with self.subTest( + op=op, + data_dtype=data_dtype, + mask_dtype=mask_dtype, + ): + specs = _select_specs( + op, + data_dtype, + mask_dtype=mask_dtype, + ) + candidates = legal_candidates(op, "a5", specs) + base_name = "template_tsel" if op == "pto.tsel" else "template_tsels" + self.assertEqual( + [candidate.name for candidate in candidates], + [f"{base_name}_1d", base_name], + ) + self.assertEqual( + [candidate.metadata.id for candidate in candidates], + [1, 0], + ) + self.assertEqual( + [candidate.metadata.loop_depth for candidate in candidates], + [1, 2], + ) + + mlir = candidates[0].specialize(**specs).mlir_text() + self.assertEqual(mlir.count("scf.for"), 1) + self.assertIn("pto.plds", mlir) + self.assertIn("pto.vsel", mlir) + self.assertIn(expected_dist, mlir) + if op == "pto.tsels": + self.assertIn("pto.vdup", mlir) + + def test_select_family_uses_predicate_specific_fallbacks(self): + cases = ( + ( + "partial data row", + _select_specs( + "pto.tsel", + "f32", + data_shape=(4, 256), + data_valid_shape=(4, 255), + ), + ), + ( + "predicate row padding", + _select_specs( + "pto.tsels", + "f32", + data_shape=(4, 256), + mask_shape=(4, 64), + ), + ), + ( + "predicate stride gap", + _select_specs( + "pto.tsels", + "f32", + data_shape=(4, 256), + mask_compact_mode=2, + ), + ), + ( + "temporary stride gap", + _select_specs( + "pto.tsel", + "f32", + data_shape=(4, 256), + tmp_compact_mode=2, + ), + ), + ) + for label, specs in cases: + with self.subTest(case=label): + op = "pto.tsels" if "scalar" in specs else "pto.tsel" + candidates = legal_candidates(op, "a5", specs) + fallback = ( + "template_tsels" if op == "pto.tsels" else "template_tsel" + ) + self.assertEqual( + [candidate.name for candidate in candidates], + [fallback], + ) + self.assertEqual(candidates[0].metadata.loop_depth, 2) + + def test_select_family_aligned_multi_row_selects_flattened_body(self): + for op in ("pto.tsel", "pto.tsels"): + with self.subTest(op=op): + specs = _select_specs( + op, + "f32", + data_shape=(4, 256), + ) + selected = select(op, "a5", specs) + mlir = selected.specialize(**specs).mlir_text() + + self.assertEqual(selected.name, f"template_{op[4:]}_1d") + self.assertEqual(selected.metadata.loop_depth, 1) + self.assertEqual(mlir.count("scf.for"), 1) + + def test_select_f32_fallback_processes_a_65_element_row_as_a_pair(self): + for op, module in ( + ("pto.tsel", tsel_templates), + ("pto.tsels", tsels_templates), + ): + with self.subTest(op=op): + self.assertEqual(module._f32_paired_cols(63, 64), 0) + self.assertEqual(module._f32_paired_cols(65, 64), 128) + self.assertEqual(module._f32_paired_cols(129, 64), 128) + + specs = _select_specs( + op, + "f32", + data_shape=(4, 128), + data_valid_shape=(4, 65), + ) + fallback = "template_tsel" if op == "pto.tsel" else "template_tsels" + selected = select(op, "a5", specs, candidate_id=fallback) + mlir = selected.specialize(**specs).mlir_text() + + self.assertIn("pto.pintlv_b16", mlir) + + def test_single_row_tcmp_mask_is_legal_input_to_tsel_1d(self): + data = TileSpec( + shape=(1, 256), + valid_shape=(1, 255), + dtype=ScalarType("f32"), + ) + mask = TileSpec( + shape=(1, 256), + valid_shape=(1, 255), + dtype=ScalarType("i8"), + ) + compare_specs = {"src0": data, "src1": data, "dst": mask} + select_specs = { + "mask": mask, + "src0": data, + "src1": data, + "tmp": TileSpec( + shape=(1, 64), + dtype=ScalarType("f32"), + ), + "dst": data, + } + + self.assertEqual( + select("pto.tcmp", "a5", compare_specs).name, + "template_tcmp_1d", + ) + self.assertEqual( + select("pto.tsel", "a5", select_specs).name, + "template_tsel_1d", + ) + + def test_tdiv_mismatched_ranges_retain_existing_2d_candidate(self): + specs = { + "src0": TileSpec(shape=(4, 65), dtype=ScalarType("f32")), + "src1": TileSpec( + shape=(4, 65), + valid_shape=(3, 65), + dtype=ScalarType("f32"), + ), + "dst": TileSpec(shape=(4, 65), dtype=ScalarType("f32")), + } + + candidates = legal_candidates("pto.tdiv", "a5", specs) + self.assertEqual( + [candidate.name for candidate in candidates], + ["template_tdiv"], + ) + self.assertEqual(candidates[0].metadata.loop_depth, 2) + if __name__ == "__main__": unittest.main() diff --git a/ptodsl/tests/test_tilelib_select.py b/ptodsl/tests/test_tilelib_select.py index 21cd4d2836..4a5edde3a8 100644 --- a/ptodsl/tests/test_tilelib_select.py +++ b/ptodsl/tests/test_tilelib_select.py @@ -11,6 +11,7 @@ from types import SimpleNamespace from ptodsl.tilelib import ( + AmbiguousTemplate, ScalarSpec, ScalarType, TemplateMetadata, @@ -40,6 +41,86 @@ def _plain_specs(*, dtype="f32", memory_space="ub", b_layout="row_major"): class TileLibSelectTest(unittest.TestCase): + def test_priority_order_is_independent_of_registration_and_candidate_id(self): + def descriptor(name, priority, candidate_id): + return SimpleNamespace( + op="pto.test_order", + target="a5", + name=name, + param_names=("src0", "src1", "dst"), + metadata=TemplateMetadata.build( + op="pto.test_order", + target="a5", + name=name, + priority=priority, + id=candidate_id, + ), + ) + + preferred = descriptor("preferred_1d", priority=10, candidate_id=99) + fallback = descriptor("fallback_2d", priority=0, candidate_id=0) + for registration_order in ( + (fallback, preferred), + (preferred, fallback), + ): + with self.subTest( + registration_order=[ + candidate.name for candidate in registration_order + ] + ): + registry = TileTemplateRegistry() + for candidate in registration_order: + registry.register(candidate) + + legal = registry.legal_candidates( + "pto.test_order", + "a5", + _f32_specs(), + ) + + self.assertEqual( + [candidate.name for candidate in legal], + ["preferred_1d", "fallback_2d"], + ) + self.assertEqual( + registry.select( + "pto.test_order", + "a5", + _f32_specs(), + ).name, + "preferred_1d", + ) + + def test_equal_top_priority_is_ambiguous_independent_of_registration_order(self): + registry = TileTemplateRegistry() + for name in ("second", "first"): + registry.register( + SimpleNamespace( + op="pto.test_tie", + target="a5", + name=name, + param_names=("src0", "src1", "dst"), + metadata=TemplateMetadata.build( + op="pto.test_tie", + target="a5", + name=name, + priority=1, + ), + ) + ) + + legal = registry.legal_candidates( + "pto.test_tie", + "a5", + _f32_specs(), + ) + self.assertEqual( + [candidate.name for candidate in legal], + ["first", "second"], + ) + with self.assertRaises(AmbiguousTemplate): + registry.select("pto.test_tie", "a5", _f32_specs()) + def test_hard_metadata_legality_is_centralized(self): registry = TileTemplateRegistry() registry.register(SimpleNamespace( @@ -148,12 +229,15 @@ def test_scalar_operand_dtypes_participate_in_legality(self): }, ) - def test_tadd_uses_single_elementwise_candidate(self): + def test_tadd_prefers_1d_and_retains_2d_fallback(self): candidates = legal_candidates("pto.tadd", "a5", _f32_specs()) - self.assertEqual([candidate.name for candidate in candidates], ["template_tadd"]) + self.assertEqual( + [candidate.name for candidate in candidates], + ["template_tadd_1d", "template_tadd"], + ) chosen = select("pto.tadd", "a5", _f32_specs()) - self.assertEqual(chosen.name, "template_tadd") + self.assertEqual(chosen.name, "template_tadd_1d") self.assertEqual( chosen.metadata.dtypes, ( @@ -169,7 +253,7 @@ def test_tadd_uses_single_elementwise_candidate(self): ), ) self.assertFalse(chosen.metadata.is_post_update) - self.assertEqual(chosen.metadata.loop_depth, 2) + self.assertEqual(chosen.metadata.loop_depth, 1) self.assertIsNone(chosen.metadata.Tail) self.assertEqual(chosen.metadata.iteration_axis, "none") self.assertEqual(chosen.metadata.op_engine, "vector") diff --git a/ptodsl/tests/test_tilelib_template_package.py b/ptodsl/tests/test_tilelib_template_package.py new file mode 100644 index 0000000000..20ed7b82ce --- /dev/null +++ b/ptodsl/tests/test_tilelib_template_package.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""Tests for TileOps package resolution outside the PTOAS CLI launcher.""" + +import unittest +from unittest.mock import Mock, patch + +from ptodsl.tilelib import _template_package + + +class TileOpsPackageResolutionTest(unittest.TestCase): + def setUp(self): + _template_package.tileops_package.cache_clear() + + def tearDown(self): + _template_package.tileops_package.cache_clear() + + def test_prefers_top_level_source_package(self): + source_package = Mock() + with patch.object( + _template_package, + "import_module", + return_value=source_package, + ) as import_module: + self.assertIs(_template_package.tileops_package(), source_package) + + import_module.assert_called_once_with("TileOps") + + def test_falls_back_to_packaged_runtime_resources(self): + packaged = Mock() + missing_source = ModuleNotFoundError( + "No module named 'TileOps'", + name="TileOps", + ) + with patch.object( + _template_package, + "import_module", + side_effect=(missing_source, packaged), + ) as import_module: + self.assertIs(_template_package.tileops_package(), packaged) + + self.assertEqual( + [call.args[0] for call in import_module.call_args_list], + ["TileOps", "ptoas._runtime.share.ptoas.TileOps"], + ) + + def test_does_not_hide_source_package_dependency_errors(self): + missing_dependency = ModuleNotFoundError( + "No module named 'dependency'", + name="dependency", + ) + with patch.object( + _template_package, + "import_module", + side_effect=missing_dependency, + ) as import_module: + with self.assertRaises(ModuleNotFoundError) as raised: + _template_package.tileops_package() + + self.assertIs(raised.exception, missing_dependency) + import_module.assert_called_once_with("TileOps") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/lit/tile_fusion/op_fusion_backend_lifecycle_level3.pto b/test/lit/tile_fusion/op_fusion_backend_lifecycle_level3.pto index 871fde9149..8ca803afa7 100644 --- a/test/lit/tile_fusion/op_fusion_backend_lifecycle_level3.pto +++ b/test/lit/tile_fusion/op_fusion_backend_lifecycle_level3.pto @@ -58,8 +58,11 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind (index, index, index) { +// LLF: %[[ROWS:.*]] = arith.constant 32 : index +// LLF: %[[COLS:.*]] = arith.constant 32 : index +// LLF: %[[ELEMENT_COUNT:.*]] = arith.muli %[[ROWS]], %[[COLS]] : index +// LLF: %[[VECTOR_WIDTH:.*]] = arith.constant 64 : index +// LLF: scf.for {{.*}} to %[[ELEMENT_COUNT]] step %[[VECTOR_WIDTH]] {{.*}} -> (index, index, index) { // LLF: pto.plt_b32 // LLF: pto.vadd // LLF: pto.vsts @@ -80,10 +83,12 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind&1 | FileCheck %s --check-prefix=META +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT + +// META-LABEL: func.func @TCMP_1D +// META: pto.tcmp +// META-SAME: candidates = [{id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tcmp_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tcmp" + +// META-LABEL: func.func @TCMP_2D +// META: pto.tcmp +// META-SAME: candidates = [{id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tcmp" +// META-NOT: name = "template_tcmp_1d" + +// META-LABEL: func.func @TCMPS_1D +// META: pto.tcmps +// META-SAME: candidates = [{id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tcmps_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tcmps" + +// META-LABEL: func.func @TCMPS_2D +// META: pto.tcmps +// META-SAME: candidates = [{id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tcmps" +// META-NOT: name = "template_tcmps_1d" + +// SELECT-LABEL: func.func private @{{.*}}__template_tcmp_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vcmp +// SELECT: pto.pdintlv_b8 +// SELECT: pto.psts +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tcmp( +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vcmp +// SELECT: pto.psts +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tcmps_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vcmps +// SELECT: pto.pdintlv_b8 +// SELECT: pto.psts +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tcmps( +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vcmps +// SELECT: pto.psts +// SELECT: return + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @TCMP_1D() { + %src0 = pto.alloc_tile + : !pto.tile_buf + %src1 = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tcmp ins( + %src0, %src1 {cmpMode = #pto} + : !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TCMP_2D() { + %src0 = pto.alloc_tile + : !pto.tile_buf + %src1 = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tcmp ins( + %src0, %src1 {cmpMode = #pto} + : !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TCMPS_1D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + %scalar = arith.constant 1.0 : f32 + + pto.tcmps ins( + %src, %scalar {cmpMode = #pto} + : !pto.tile_buf, + f32) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TCMPS_2D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + %scalar = arith.constant 1.0 : f32 + + pto.tcmps ins( + %src, %scalar {cmpMode = #pto} + : !pto.tile_buf, + f32) + outs( + %dst + : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/expand_tile_op_ptodsl_conversion_1d_2d.pto b/test/lit/vpto/expand_tile_op_ptodsl_conversion_1d_2d.pto new file mode 100644 index 0000000000..6c847b7ea9 --- /dev/null +++ b/test/lit/vpto/expand_tile_op_ptodsl_conversion_1d_2d.pto @@ -0,0 +1,91 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// Width-changing conversion preserves ranked 1D/2D selection through +// InsertTemplateAttributes and ExpandTileOp. The contiguous full-axis case +// uses one flattened loop; a multi-row partial column axis keeps the row-wise +// fallback and its two loops. +// +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT + +// META-LABEL: func.func @TCVT_F32_I16_1D +// META: pto.tcvt +// META-SAME: candidates = [{ +// META-SAME: id = 53 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tcvt_f32_to_i16_1d" +// META-SAME: }, { +// META-SAME: id = 15 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tcvt_f32_to_i16" + +// META-LABEL: func.func @TCVT_F32_I16_2D +// META: pto.tcvt +// META-SAME: candidates = [{ +// META-SAME: id = 15 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tcvt_f32_to_i16" +// META-NOT: name = "template_tcvt_f32_to_i16_1d" + +// SELECT-LABEL: func.func private @{{.*}}__template_tcvt_f32_to_i16_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vcvt +// SELECT: pto.vcvt +// SELECT: pto.vsts {{.*}} {dist = "PK_B32"} +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tcvt_f32_to_i16( +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vcvt +// SELECT: pto.vcvt +// SELECT: pto.vsts {{.*}} {dist = "PK_B32"} +// SELECT: return + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @TCVT_F32_I16_1D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tcvt ins( + %src + : !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TCVT_F32_I16_2D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tcvt ins( + %src + : !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/expand_tile_op_ptodsl_scalar_1d_2d.pto b/test/lit/vpto/expand_tile_op_ptodsl_scalar_1d_2d.pto new file mode 100644 index 0000000000..815a54beda --- /dev/null +++ b/test/lit/vpto/expand_tile_op_ptodsl_scalar_1d_2d.pto @@ -0,0 +1,89 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// Tile-Scalar candidates apply flattened legality to the source and +// destination tiles. The scalar operand has no layout or continuity metadata. +// +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT + +// META-LABEL: func.func @TADDS_1D +// META: pto.tadds +// META-SAME: candidates = [{ +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tadds_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tadds" + +// META-LABEL: func.func @TADDS_2D +// META: pto.tadds +// META-SAME: candidates = [{ +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tadds" +// META-NOT: name = "template_tadds_1d" + +// SELECT-LABEL: func.func private @{{.*}}__template_tadds_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vadds +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tadds( +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vadds +// SELECT: return + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @TADDS_1D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + %scalar = arith.constant 1.0 : f32 + + pto.tadds ins( + %src, %scalar + : !pto.tile_buf, + f32) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TADDS_2D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + %scalar = arith.constant 1.0 : f32 + + pto.tadds ins( + %src, %scalar + : !pto.tile_buf, + f32) + outs( + %dst + : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/expand_tile_op_ptodsl_scalar_fill_1d_2d.pto b/test/lit/vpto/expand_tile_op_ptodsl_scalar_fill_1d_2d.pto new file mode 100644 index 0000000000..403cf10ffb --- /dev/null +++ b/test/lit/vpto/expand_tile_op_ptodsl_scalar_fill_1d_2d.pto @@ -0,0 +1,77 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// Scalar-fill flattened legality depends only on the destination tile. The +// scalar operand has no memory layout or continuity requirement. +// +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT + +// META-LABEL: func.func @TEXPANDS_1D +// META: pto.texpands +// META-SAME: candidates = [{ +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_texpands_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_texpands" + +// META-LABEL: func.func @TEXPANDS_2D +// META: pto.texpands +// META-SAME: candidates = [{ +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_texpands" +// META-NOT: name = "template_texpands_1d" + +// SELECT-LABEL: func.func private @{{.*}}__template_texpands_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vdup +// SELECT: pto.vsts +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_texpands( +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vdup +// SELECT: pto.vsts +// SELECT: return + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @TEXPANDS_1D() { + %scalar = arith.constant 1.0 : f32 + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.texpands ins(%scalar : f32) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TEXPANDS_2D() { + %scalar = arith.constant 1.0 : f32 + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.texpands ins(%scalar : f32) + outs( + %dst + : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/expand_tile_op_ptodsl_select_1d_2d.pto b/test/lit/vpto/expand_tile_op_ptodsl_select_1d_2d.pto new file mode 100644 index 0000000000..f0b3e110c6 --- /dev/null +++ b/test/lit/vpto/expand_tile_op_ptodsl_select_1d_2d.pto @@ -0,0 +1,221 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// Predicate select may flatten one sufficiently padded logical row or aligned +// multi-row data with an exactly dense packed-mask row stride. +// +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT + +// META-LABEL: func.func @TSEL_1D +// META: pto.tsel +// META-SAME: candidates = [{ +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tsel_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tsel" + +// META-LABEL: func.func @TSEL_2D +// META: pto.tsel +// META-SAME: candidates = [{ +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tsel" +// META-NOT: name = "template_tsel_1d" + +// META-LABEL: func.func @TSELS_1D +// META: pto.tsels +// META-SAME: candidates = [{ +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tsels_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tsels" + +// META-LABEL: func.func @TSELS_2D +// META: pto.tsels +// META-SAME: candidates = [{ +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tsels" +// META-NOT: name = "template_tsels_1d" + +// SELECT-LABEL: func.func private @{{.*}}__template_tsel_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.plds +// SELECT: pto.pintlv_b16 +// SELECT: pto.vsel +// SELECT: pto.vsts +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tsel( +// SELECT: scf.for +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.plds +// SELECT: pto.vsel +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tsels_1d( +// SELECT: pto.vdup +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.plds +// SELECT: pto.pintlv_b16 +// SELECT: pto.vsel +// SELECT: pto.vsts +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tsels( +// SELECT: pto.vdup +// SELECT: scf.for +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.plds +// SELECT: pto.vsel +// SELECT: return + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @TSEL_1D() { + %mask = pto.alloc_tile + : !pto.tile_buf + %src0 = pto.alloc_tile + : !pto.tile_buf + %src1 = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tsel ins( + %mask, %src0, %src1, %tmp + : !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TSEL_2D() { + %mask = pto.alloc_tile + : !pto.tile_buf + %src0 = pto.alloc_tile + : !pto.tile_buf + %src1 = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tsel ins( + %mask, %src0, %src1, %tmp + : !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TSELS_1D() { + %mask = pto.alloc_tile + : !pto.tile_buf + %src = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + %scalar = arith.constant 1.0 : f32 + + pto.tsels ins( + %mask, %src, %tmp, %scalar + : !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf, + f32) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TSELS_2D() { + %mask = pto.alloc_tile + : !pto.tile_buf + %src = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + %scalar = arith.constant 1.0 : f32 + + pto.tsels ins( + %mask, %src, %tmp, %scalar + : !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf, + f32) + outs( + %dst + : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/expand_tile_op_ptodsl_specialized_binary_1d_2d.pto b/test/lit/vpto/expand_tile_op_ptodsl_specialized_binary_1d_2d.pto new file mode 100644 index 0000000000..81ee8fad3c --- /dev/null +++ b/test/lit/vpto/expand_tile_op_ptodsl_specialized_binary_1d_2d.pto @@ -0,0 +1,187 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// Specialized Tile-Tile algorithms stay in their own modules while using the +// shared traversal forms. High-precision division remains specialized by +// precisionType, and fmod retains its remainder instruction sequence. +// +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT + +// META-LABEL: func.func @TDIV_HIGH_1D +// META: pto.tdiv +// META-SAME: candidates = [{ +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tdiv_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tdiv" + +// META-LABEL: func.func @TDIV_HIGH_2D +// META: pto.tdiv +// META-SAME: candidates = [{ +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tdiv" +// META-NOT: name = "template_tdiv_1d" + +// META-LABEL: func.func @TFMOD_1D +// META: pto.tfmod +// META-SAME: candidates = [{ +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tfmod_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tfmod" + +// META-LABEL: func.func @TFMOD_2D +// META: pto.tfmod +// META-SAME: candidates = [{ +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tfmod" +// META-NOT: name = "template_tfmod_1d" + +// SELECT-LABEL: func.func private @{{.*}}__template_tdiv_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vdiv +// SELECT: pto.vbitcast +// SELECT: pto.vcmp +// SELECT: pto.vsel +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tdiv( +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vdiv +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tfmod_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vdiv +// SELECT: pto.vtrc +// SELECT: pto.vmul +// SELECT: pto.vsub +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tfmod( +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vdiv +// SELECT: pto.vtrc +// SELECT: pto.vmul +// SELECT: pto.vsub +// SELECT: return + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @TDIV_HIGH_1D() { + %lhs = pto.alloc_tile + : !pto.tile_buf + %rhs = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tdiv ins( + %lhs, %rhs + : !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + {precisionType = #pto} + return + } + + func.func @TDIV_HIGH_2D() { + %lhs = pto.alloc_tile + : !pto.tile_buf + %rhs = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tdiv ins( + %lhs, %rhs + : !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + {precisionType = #pto} + return + } + + func.func @TFMOD_1D() { + %lhs = pto.alloc_tile + : !pto.tile_buf + %rhs = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tfmod ins( + %lhs, %rhs + : !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TFMOD_2D() { + %lhs = pto.alloc_tile + : !pto.tile_buf + %rhs = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tfmod ins( + %lhs, %rhs + : !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/expand_tile_op_ptodsl_specialized_scalar_1d_2d.pto b/test/lit/vpto/expand_tile_op_ptodsl_specialized_scalar_1d_2d.pto new file mode 100644 index 0000000000..718e73a83e --- /dev/null +++ b/test/lit/vpto/expand_tile_op_ptodsl_specialized_scalar_1d_2d.pto @@ -0,0 +1,187 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// Specialized Tile-Scalar algorithms retain operand order, precision handling, +// and temporary operands while sharing flattened and row-wise traversal. +// +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT + +// META-LABEL: func.func @TDIVS_TILE_1D +// META: pto.tdivs +// META-SAME: candidates = [{ +// META-SAME: id = 2 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tdivs_tile_scalar_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tdivs_tile_scalar" + +// META-LABEL: func.func @TDIVS_SCALAR_1D +// META: pto.tdivs +// META-SAME: candidates = [{ +// META-SAME: id = 3 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tdivs_scalar_tile_1d" +// META-SAME: }, { +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tdivs_scalar_tile" + +// META-LABEL: func.func @TREMS_1D +// META: pto.trems +// META-SAME: candidates = [{ +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_trems_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_trems" + +// META-LABEL: func.func @TREMS_TMP_2D +// META: pto.trems +// META-SAME: candidates = [{ +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_trems" +// META-NOT: name = "template_trems_1d" + +// SELECT-LABEL: func.func private @{{.*}}__template_tdivs_tile_scalar_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vbr +// SELECT: pto.vdiv +// SELECT: pto.vbitcast +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tdivs_scalar_tile_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vbr +// SELECT: pto.vdiv +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_trems_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vdiv +// SELECT: pto.vtrc +// SELECT: pto.vmuls +// SELECT: pto.vsub +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_trems( +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vdiv +// SELECT: pto.vtrc +// SELECT: pto.vmuls +// SELECT: pto.vsub +// SELECT: return + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @TDIVS_TILE_1D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + %scalar = arith.constant 2.0 : f32 + + pto.tdivs ins( + %src, %scalar + : !pto.tile_buf, + f32) + outs( + %dst + : !pto.tile_buf) + {precisionType = #pto} + return + } + + func.func @TDIVS_SCALAR_1D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + %scalar = arith.constant 2.0 : f32 + + pto.tdivs ins( + %scalar, %src + : f32, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TREMS_1D() { + %src = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + %scalar = arith.constant 3.0 : f32 + + pto.trems ins( + %src, %scalar, %tmp + : !pto.tile_buf, + f32, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TREMS_TMP_2D() { + %src = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + %scalar = arith.constant 3.0 : f32 + + pto.trems ins( + %src, %scalar, %tmp + : !pto.tile_buf, + f32, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/expand_tile_op_ptodsl_specialized_unary_1d_2d.pto b/test/lit/vpto/expand_tile_op_ptodsl_specialized_unary_1d_2d.pto new file mode 100644 index 0000000000..e26866a34f --- /dev/null +++ b/test/lit/vpto/expand_tile_op_ptodsl_specialized_unary_1d_2d.pto @@ -0,0 +1,131 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// Specialized unary algorithms remain in their operation modules while using +// the common traversal forms. Precision filtering happens before traversal +// ranking, so high-precision tlog never competes with its default candidates. +// +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT + +// META-LABEL: func.func @TLOG_HIGH_1D +// META: pto.tlog +// META-SAME: candidates = [{ +// META-SAME: id = 3 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tlog_high_precision_1d" +// META-SAME: }, { +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tlog_high_precision" +// META-NOT: name = "template_tlog_1d" + +// META-LABEL: func.func @TLOG_HIGH_2D +// META: pto.tlog +// META-SAME: candidates = [{ +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tlog_high_precision" +// META-NOT: name = "template_tlog_high_precision_1d" + +// META-LABEL: func.func @TRECIP_1D +// META: pto.trecip +// META-SAME: candidates = [{ +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_trecip_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_trecip" + +// SELECT-LABEL: func.func private @{{.*}}__template_tlog_high_precision_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vcmps +// SELECT: pto.vmuls +// SELECT: pto.vsel +// SELECT: pto.vln +// SELECT: pto.vadds +// SELECT: pto.vsel +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tlog_high_precision( +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vcmps +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_trecip_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vbr +// SELECT: pto.vdiv +// SELECT: return + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @TLOG_HIGH_1D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tlog ins( + %src + : !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + {precisionType = #pto} + return + } + + func.func @TLOG_HIGH_2D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tlog ins( + %src + : !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + {precisionType = #pto} + return + } + + func.func @TRECIP_1D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.trecip ins( + %src + : !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/expand_tile_op_ptodsl_tabs_1d_2d.pto b/test/lit/vpto/expand_tile_op_ptodsl_tabs_1d_2d.pto new file mode 100644 index 0000000000..71f8883d6b --- /dev/null +++ b/test/lit/vpto/expand_tile_op_ptodsl_tabs_1d_2d.pto @@ -0,0 +1,84 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// A contiguous full-axis unary operation records the preferred 1D candidate +// first and expands to its one-loop helper. A multi-row partial-column +// operation records only the existing 2D fallback and expands to two loops. +// +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT + +// META-LABEL: func.func @TABS_1D +// META: pto.tabs +// META-SAME: candidates = [{ +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tabs_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tabs" + +// META-LABEL: func.func @TABS_2D +// META: pto.tabs +// META-SAME: candidates = [{ +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tabs" +// META-NOT: name = "template_tabs_1d" + +// SELECT-LABEL: func.func private @{{.*}}__template_tabs_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tabs( +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: return + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @TABS_1D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tabs ins( + %src + : !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TABS_2D() { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tabs ins( + %src + : !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto b/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto index fbef117298..48cd9749c0 100644 --- a/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto +++ b/test/lit/vpto/expand_tile_op_ptodsl_tadd.pto @@ -18,23 +18,46 @@ // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>/dev/null | FileCheck %s --check-prefix=EXPAND +// META-LABEL: func.func @TADD // META: pto.tadd -// META-SAME: candidates = [ +// META-SAME: candidates = [{ +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tadd_1d" +// META-SAME: postupdate = 0 : i64 +// META-SAME: tail = 0 : i64}, { // META-SAME: id = 0 : i64 // META-SAME: loop_depth = 2 : i64 // META-SAME: name = "template_tadd" -// META-SAME: postupdate = 0 : i64 -// META-SAME: tail = 0 : i64}] // META-NOT: priority = // META-NOT: tags = // META-NOT: fusible = +// META-LABEL: func.func @TADD_2D +// META: pto.tadd +// META-SAME: candidates = [{ +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tadd" +// META-NOT: name = "template_tadd_1d" + // PREFUSION: IR Dump Before FusionPlan // PREFUSION: pto.tadd // PREFUSION-SAME: candidates = [ -// SELECT-COUNT-2: call @{{.*}}__template_tadd -// SELECT-COUNT-1: func.func {{.*}}@{{.*}}__template_tadd +// SELECT-COUNT-2: call @{{.*}}__template_tadd_1d( +// SELECT-COUNT-1: call @{{.*}}__template_tadd( + +// SELECT-COUNT-1: func.func private @{{.*}}__template_tadd_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: return + +// SELECT-COUNT-1: func.func private @{{.*}}__template_tadd( +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: return // EXPAND: func.func @TADD // EXPAND-NOT: pto.tadd ins @@ -82,4 +105,28 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { blayout=row_major, slayout=none_box, fractal=512, pad=0>) return } + + func.func @TADD_2D() { + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tadd ins( + %a, %b + : !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } } diff --git a/test/lit/vpto/expand_tile_op_ptodsl_temporary_binary_1d_2d.pto b/test/lit/vpto/expand_tile_op_ptodsl_temporary_binary_1d_2d.pto new file mode 100644 index 0000000000..415cf2e609 --- /dev/null +++ b/test/lit/vpto/expand_tile_op_ptodsl_temporary_binary_1d_2d.pto @@ -0,0 +1,109 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// A temporary TileOp operand participates in flattened-traversal legality even +// when the generated computation does not access it. The first PReLU has four +// contiguous logical ranges and prefers the 1D helper. In the second, only the +// i8 temporary has a row-plus-one stride, so only the 2D fallback is legal. +// +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=META +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELECT + +// META-LABEL: func.func @TPRELU_1D +// META: pto.tprelu +// META-SAME: candidates = [{ +// META-SAME: id = 1 : i64 +// META-SAME: loop_depth = 1 : i64 +// META-SAME: name = "template_tprelu_1d" +// META-SAME: }, { +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tprelu" + +// META-LABEL: func.func @TPRELU_TMP_2D +// META: pto.tprelu +// META-SAME: candidates = [{ +// META-SAME: id = 0 : i64 +// META-SAME: loop_depth = 2 : i64 +// META-SAME: name = "template_tprelu" +// META-NOT: name = "template_tprelu_1d" + +// SELECT-LABEL: func.func private @{{.*}}__template_tprelu_1d( +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vprelu +// SELECT: return + +// SELECT-LABEL: func.func private @{{.*}}__template_tprelu( +// SELECT: scf.for +// SELECT: scf.for +// SELECT-NOT: scf.for +// SELECT: pto.vprelu +// SELECT: return + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @TPRELU_1D() { + %src0 = pto.alloc_tile + : !pto.tile_buf + %src1 = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tprelu ins( + %src0, %src1, %tmp + : !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } + + func.func @TPRELU_TMP_2D() { + %src0 = pto.alloc_tile + : !pto.tile_buf + %src1 = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tprelu ins( + %src0, %src1, %tmp + : !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/expand_tile_op_tilelang_tdivs.pto b/test/lit/vpto/expand_tile_op_tilelang_tdivs.pto index 49148165aa..7be9c5cff8 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tdivs.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tdivs.pto @@ -10,9 +10,9 @@ // CHECK-TILE-SCALAR-NOT: pto.tdivs ins // CHECK-TILE-SCALAR: pto.vecscope // CHECK-TILE-SCALAR: pto.castptr -// CHECK-TILE-SCALAR: %[[MASK:.+]], %[[SCALAR_OUT:.+]] = pto.plt_b32 // CHECK-TILE-SCALAR: %[[BR:.+]] = pto.vbr // CHECK-TILE-SCALAR: scf.for +// CHECK-TILE-SCALAR: %[[MASK:.+]], %[[SCALAR_OUT:.+]] = pto.plt_b32 // CHECK-TILE-SCALAR: %[[LD:.+]] = pto.vlds // CHECK-TILE-SCALAR: %[[DIV:.+]] = pto.vdiv %[[LD]], %[[BR]], %[[MASK]] // CHECK-TILE-SCALAR: pto.vsts %[[DIV]] @@ -22,9 +22,9 @@ // CHECK-SCALAR-TILE-NOT: pto.tdivs ins // CHECK-SCALAR-TILE: pto.vecscope // CHECK-SCALAR-TILE: pto.castptr -// CHECK-SCALAR-TILE: %[[MASK:.+]], %[[SCALAR_OUT:.+]] = pto.plt_b32 // CHECK-SCALAR-TILE: %[[BR:.+]] = pto.vbr // CHECK-SCALAR-TILE: scf.for +// CHECK-SCALAR-TILE: %[[MASK:.+]], %[[SCALAR_OUT:.+]] = pto.plt_b32 // CHECK-SCALAR-TILE: %[[LD:.+]] = pto.vlds // CHECK-SCALAR-TILE: %[[DIV:.+]] = pto.vdiv %[[BR]], %[[LD]], %[[MASK]] // CHECK-SCALAR-TILE: pto.vsts %[[DIV]] diff --git a/test/lit/vpto/fold_tile_buf_intrinsics.pto b/test/lit/vpto/fold_tile_buf_intrinsics.pto index 7b23584bca..8f872a449c 100644 --- a/test/lit/vpto/fold_tile_buf_intrinsics.pto +++ b/test/lit/vpto/fold_tile_buf_intrinsics.pto @@ -29,14 +29,15 @@ // - tile_buf_addr should still carry the planned tile address as pto.ptr // ADDR-LABEL: func.func @TADD // ADDR: pto.tile_buf_addr {{.*}} -> {{.*}} -// ADDR: pto.vlds // ADDR: pto.tile_buf_addr {{.*}} -> {{.*}} -// ADDR: pto.vlds -// ADDR: pto.vadd // ADDR: pto.tile_buf_addr {{.*}} -> {{.*}} // ADDR-NOT: pto.pointer_cast // ADDR-NOT: pto.tile_valid_rows // ADDR-NOT: pto.tile_valid_cols +// ADDR: scf.for +// ADDR: pto.vlds +// ADDR: pto.vlds +// ADDR: pto.vadd // ADDR: pto.vsts // After VPTO pointer normalization: @@ -44,14 +45,15 @@ // - tile-slice addressing is carried by the vlds/vsts offset operand // NORMALIZED-LABEL: func.func @TADD // NORMALIZED: pto.castptr -// NORMALIZED: arith.muli -// NORMALIZED: pto.vlds // NORMALIZED: pto.castptr -// NORMALIZED: pto.vlds -// NORMALIZED: pto.vadd // NORMALIZED: pto.castptr // NORMALIZED-NOT: pto.tile_buf_addr // NORMALIZED-NOT: pto.pointer_cast +// NORMALIZED-NOT: arith.muli +// NORMALIZED: scf.for +// NORMALIZED: pto.vlds +// NORMALIZED: pto.vlds +// NORMALIZED: pto.vadd // NORMALIZED: pto.vsts module attributes {pto.kernel_kind = #pto.kernel_kind} { diff --git a/test/lit/vpto/insert_template_attributes_candidate_order.pto b/test/lit/vpto/insert_template_attributes_candidate_order.pto new file mode 100644 index 0000000000..058823a6d4 --- /dev/null +++ b/test/lit/vpto/insert_template_attributes_candidate_order.pto @@ -0,0 +1,62 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// The preferred in-process candidate has both higher priority and a larger ID, +// proving that candidate ID does not control compact-attribute ordering or +// expansion. +// +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=EXPAND --implicit-check-not=pto.tadd + +// CHECK: pto.tadd +// CHECK-SAME: candidates = [{ +// CHECK-SAME: id = 1 : i64 +// CHECK-SAME: loop_depth = 1 : i64 +// CHECK-SAME: name = "template_tadd_1d" +// CHECK-SAME: postupdate = 0 : i64 +// CHECK-SAME: tail = 0 : i64 +// CHECK-SAME: }, { +// CHECK-SAME: id = 0 : i64 +// CHECK-SAME: loop_depth = 2 : i64 +// CHECK-SAME: name = "template_tadd" + +// EXPAND: func.func @candidate_order +// EXPAND: call @{{.*}}__template_tadd_1d +// EXPAND: func.func private @{{.*}}__template_tadd_1d + +module attributes {pto.target_arch = "a5"} { + func.func @candidate_order() { + %src0 = pto.alloc_tile + : !pto.tile_buf + %src1 = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tadd ins( + %src0, %src1 + : !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } +} diff --git a/test/vpto/cases/elementwise-1d-2d-equivalence.py b/test/vpto/cases/elementwise-1d-2d-equivalence.py new file mode 100644 index 0000000000..42e610701d --- /dev/null +++ b/test/vpto/cases/elementwise-1d-2d-equivalence.py @@ -0,0 +1,587 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Runtime equivalence checks through element-wise 1D/2D TileOp selection.""" + +from pathlib import Path +import sys + +import numpy as np + + +def _bootstrap_dsl_st_common() -> None: + here = Path(__file__).resolve() + for candidate in here.parents: + common_dir = candidate / "test" / "dsl-st" + if (common_dir / "common.py").exists(): + sys.path.insert(0, str(common_dir)) + return + raise RuntimeError("Unable to locate test/dsl-st/common.py") + + +_bootstrap_dsl_st_common() + +from common import assert_close, auto_main +from ptodsl import pto +from ptodsl.tilelib.templates.a5.tabs import ( + template_tabs, + template_tabs_1d, +) +from ptodsl.tilelib.templates.a5.tadd import ( + template_tadd, + template_tadd_1d, +) +from ptodsl.tilelib.templates.a5.tadds import ( + template_tadds, + template_tadds_1d, +) +from ptodsl.tilelib.templates.a5.tcmps import ( + template_tcmps, + template_tcmps_1d, +) +from ptodsl.tilelib.templates.a5.tcvt import ( + template_tcvt_f32_to_i16, + template_tcvt_f32_to_i16_1d, +) +from ptodsl.tilelib.templates.a5.texpand import ( + template_texpands, + template_texpands_1d, +) +from ptodsl.tilelib.templates.a5.tsels import ( + template_tsels, + template_tsels_1d, +) + + +SEED = 20260802 +GUARD_ELEMENTS = 32 + +TABS_SHAPE = (2, 72) +TABS_PADDED_SHAPE = (2, 80) +TADD_SHAPE = (4, 80) +TADD_PADDED_SHAPE = (4, 96) +TADDS_SHAPE = (4, 96) +TADDS_PADDED_SHAPE = (4, 128) +TCMPS_DATA_SHAPE = (4, 256) +TCMPS_MASK_SHAPE = (4, 32) +TCMPS_PADDED_DATA_SHAPE = (4, 512) +TSELS_DATA_SHAPE = (4, 256) +TSELS_MASK_SHAPE = (4, 32) +TSELS_PADDED_DATA_SHAPE = (4, 512) +TSELS_PADDED_MASK_SHAPE = (4, 64) +TCVT_SHAPE = (4, 80) +TCVT_PADDED_SHAPE = (4, 96) +TEXPANDS_SHAPE = (4, 72) +TEXPANDS_PADDED_SHAPE = (4, 80) + +# Each compact shape fills its physical column axis and selects 1D. The +# corresponding padded shape keeps the same valid region but introduces a row +# stride gap, which conservatively selects the 2D fallback. + + +_TEMPLATE_PAIRS = ( + (template_tabs_1d, template_tabs), + (template_tadd_1d, template_tadd), + (template_tadds_1d, template_tadds), + (template_tcmps_1d, template_tcmps), + (template_tsels_1d, template_tsels), + (template_tcvt_f32_to_i16_1d, template_tcvt_f32_to_i16), + (template_texpands_1d, template_texpands), +) + +for _flattened, _rowwise in _TEMPLATE_PAIRS: + if _flattened.metadata.loop_depth != 1: + raise AssertionError(f"{_flattened.name} is not a 1D candidate") + if _rowwise.metadata.loop_depth != 2: + raise AssertionError(f"{_rowwise.name} is not a 2D candidate") + + +def _row_major_view(ptr, rows, cols, *, offset=0): + if offset: + ptr = pto.addptr(ptr, offset) + return pto.make_tensor_view( + ptr, + shape=[rows, cols], + strides=[cols, 1], + ) + + +def _equivalence_jit(op_name): + return pto.jit( + name=f"elementwise_{op_name}_1d_2d_equivalence", + target="a5", + kernel_kind="vector", + mode="explicit", + insert_sync=True, + ) + + +# Only cases with the same one-input ABI and tile topology use this factory. +# Packed predicates, temporaries, and other distinct operand forms stay explicit. +def _make_single_input_equivalence_kernel( + op_name, + *, + shape, + padded_shape, + src_dtype, + dst_dtype, + operation, +): + @_equivalence_jit(op_name) + def _kernel( + src_ptr: pto.ptr(src_dtype, "gm"), + out_1d_ptr: pto.ptr(dst_dtype, "gm"), + out_2d_ptr: pto.ptr(dst_dtype, "gm"), + ): + rows, cols = shape + src_view = _row_major_view(src_ptr, rows, cols) + out_1d_view = _row_major_view( + out_1d_ptr, rows, cols, offset=GUARD_ELEMENTS + ) + out_2d_view = _row_major_view( + out_2d_ptr, rows, cols, offset=GUARD_ELEMENTS + ) + + src_1d = pto.alloc_tile( + shape=[rows, cols], dtype=src_dtype, addr=0 + ) + src_2d = pto.alloc_tile( + shape=list(padded_shape), + dtype=src_dtype, + valid_shape=[rows, cols], + addr=4096, + ) + out_1d = pto.alloc_tile( + shape=[rows, cols], dtype=dst_dtype, addr=8192 + ) + out_2d = pto.alloc_tile( + shape=list(padded_shape), + dtype=dst_dtype, + valid_shape=[rows, cols], + addr=12288, + ) + + pto.tile.load(src_view, src_1d) + pto.tile.load(src_view, src_2d) + operation(src_1d, out_1d) + operation(src_2d, out_2d) + pto.tile.store(out_1d, out_1d_view) + pto.tile.store(out_2d, out_2d_view) + + return _kernel + + +def _tabs_operation(src, dst): + pto.tile.abs(src, dst) + + +def _tadds_operation(src, dst): + pto.tile.adds(src, 7, dst) + + +def _tcvt_operation(src, dst): + pto.tile.cvt(src, dst) + + +elementwise_tabs_1d_2d_equivalence = _make_single_input_equivalence_kernel( + "tabs", + shape=TABS_SHAPE, + padded_shape=TABS_PADDED_SHAPE, + src_dtype=pto.f32, + dst_dtype=pto.f32, + operation=_tabs_operation, +) + + +@_equivalence_jit("tadd") +def elementwise_tadd_1d_2d_equivalence( + lhs_ptr: pto.ptr(pto.i16, "gm"), + rhs_ptr: pto.ptr(pto.i16, "gm"), + out_1d_ptr: pto.ptr(pto.i16, "gm"), + out_2d_ptr: pto.ptr(pto.i16, "gm"), +): + rows, cols = TADD_SHAPE + + lhs_view = _row_major_view(lhs_ptr, rows, cols) + rhs_view = _row_major_view(rhs_ptr, rows, cols) + out_1d_view = _row_major_view( + out_1d_ptr, rows, cols, offset=GUARD_ELEMENTS + ) + out_2d_view = _row_major_view( + out_2d_ptr, rows, cols, offset=GUARD_ELEMENTS + ) + + lhs_1d = pto.alloc_tile(shape=[rows, cols], dtype=pto.i16, addr=0) + rhs_1d = pto.alloc_tile(shape=[rows, cols], dtype=pto.i16, addr=4096) + lhs_2d = pto.alloc_tile( + shape=list(TADD_PADDED_SHAPE), + dtype=pto.i16, + valid_shape=[rows, cols], + addr=8192, + ) + rhs_2d = pto.alloc_tile( + shape=list(TADD_PADDED_SHAPE), + dtype=pto.i16, + valid_shape=[rows, cols], + addr=12288, + ) + out_1d = pto.alloc_tile(shape=[rows, cols], dtype=pto.i16, addr=16384) + out_2d = pto.alloc_tile( + shape=list(TADD_PADDED_SHAPE), + dtype=pto.i16, + valid_shape=[rows, cols], + addr=20480, + ) + + pto.tile.load(lhs_view, lhs_1d) + pto.tile.load(rhs_view, rhs_1d) + pto.tile.load(lhs_view, lhs_2d) + pto.tile.load(rhs_view, rhs_2d) + pto.tile.add(lhs_1d, rhs_1d, out_1d) + pto.tile.add(lhs_2d, rhs_2d, out_2d) + pto.tile.store(out_1d, out_1d_view) + pto.tile.store(out_2d, out_2d_view) + + +elementwise_tadds_1d_2d_equivalence = _make_single_input_equivalence_kernel( + "tadds", + shape=TADDS_SHAPE, + padded_shape=TADDS_PADDED_SHAPE, + src_dtype=pto.i8, + dst_dtype=pto.i8, + operation=_tadds_operation, +) + + +@_equivalence_jit("tcmps") +def elementwise_tcmps_1d_2d_equivalence( + src_ptr: pto.ptr(pto.i8, "gm"), + out_1d_ptr: pto.ptr(pto.ui8, "gm"), + out_2d_ptr: pto.ptr(pto.ui8, "gm"), +): + data_rows, data_cols = TCMPS_DATA_SHAPE + mask_rows, mask_cols = TCMPS_MASK_SHAPE + + src_view = _row_major_view(src_ptr, data_rows, data_cols) + out_1d_view = _row_major_view( + out_1d_ptr, mask_rows, mask_cols, offset=GUARD_ELEMENTS + ) + out_2d_view = _row_major_view( + out_2d_ptr, mask_rows, mask_cols, offset=GUARD_ELEMENTS + ) + + src_1d = pto.alloc_tile( + shape=[data_rows, data_cols], dtype=pto.i8, addr=0 + ) + src_2d = pto.alloc_tile( + shape=list(TCMPS_PADDED_DATA_SHAPE), + dtype=pto.i8, + valid_shape=[data_rows, data_cols], + addr=4096, + ) + out_1d = pto.alloc_tile( + shape=[mask_rows, mask_cols], + dtype=pto.ui8, + addr=8192, + ) + out_2d = pto.alloc_tile( + shape=[mask_rows, mask_cols], + dtype=pto.ui8, + addr=12288, + ) + + pto.tile.load(src_view, src_1d) + pto.tile.load(src_view, src_2d) + pto.tile.cmps(src_1d, 5, out_1d) + pto.tile.cmps(src_2d, 5, out_2d) + pto.tile.store(out_1d, out_1d_view) + pto.tile.store(out_2d, out_2d_view) + + +@_equivalence_jit("tsels") +def elementwise_tsels_1d_2d_equivalence( + mask_ptr: pto.ptr(pto.i8, "gm"), + src_ptr: pto.ptr(pto.i8, "gm"), + out_1d_ptr: pto.ptr(pto.i8, "gm"), + out_2d_ptr: pto.ptr(pto.i8, "gm"), +): + data_rows, data_cols = TSELS_DATA_SHAPE + mask_rows, mask_cols = TSELS_MASK_SHAPE + + mask_view = _row_major_view(mask_ptr, mask_rows, mask_cols) + src_view = _row_major_view(src_ptr, data_rows, data_cols) + out_1d_view = _row_major_view( + out_1d_ptr, data_rows, data_cols, offset=GUARD_ELEMENTS + ) + out_2d_view = _row_major_view( + out_2d_ptr, data_rows, data_cols, offset=GUARD_ELEMENTS + ) + + mask_1d = pto.alloc_tile( + shape=[mask_rows, mask_cols], dtype=pto.i8, addr=0 + ) + src_1d = pto.alloc_tile( + shape=[data_rows, data_cols], + dtype=pto.i8, + addr=4096, + ) + tmp_1d = pto.alloc_tile(shape=[1, 256], dtype=pto.i8, addr=8192) + out_1d = pto.alloc_tile( + shape=[data_rows, data_cols], + dtype=pto.i8, + addr=12288, + ) + mask_2d = pto.alloc_tile( + shape=list(TSELS_PADDED_MASK_SHAPE), + dtype=pto.i8, + valid_shape=[mask_rows, mask_cols], + addr=16384, + ) + src_2d = pto.alloc_tile( + shape=list(TSELS_PADDED_DATA_SHAPE), + dtype=pto.i8, + valid_shape=[data_rows, data_cols], + addr=20480, + ) + tmp_2d = pto.alloc_tile(shape=[1, 256], dtype=pto.i8, addr=24576) + out_2d = pto.alloc_tile( + shape=list(TSELS_PADDED_DATA_SHAPE), + dtype=pto.i8, + valid_shape=[data_rows, data_cols], + addr=28672, + ) + + pto.tile.load(mask_view, mask_1d) + pto.tile.load(src_view, src_1d) + pto.tile.load(mask_view, mask_2d) + pto.tile.load(src_view, src_2d) + pto.tile.sels(mask_1d, src_1d, -3, out_1d, tmp=tmp_1d) + pto.tile.sels(mask_2d, src_2d, -3, out_2d, tmp=tmp_2d) + pto.tile.store(out_1d, out_1d_view) + pto.tile.store(out_2d, out_2d_view) + + +elementwise_tcvt_1d_2d_equivalence = _make_single_input_equivalence_kernel( + "tcvt", + shape=TCVT_SHAPE, + padded_shape=TCVT_PADDED_SHAPE, + src_dtype=pto.f32, + dst_dtype=pto.i16, + operation=_tcvt_operation, +) + + +@_equivalence_jit("texpands") +def elementwise_texpands_1d_2d_equivalence( + out_1d_ptr: pto.ptr(pto.i32, "gm"), + out_2d_ptr: pto.ptr(pto.i32, "gm"), +): + rows, cols = TEXPANDS_SHAPE + + out_1d_view = _row_major_view( + out_1d_ptr, rows, cols, offset=GUARD_ELEMENTS + ) + out_2d_view = _row_major_view( + out_2d_ptr, rows, cols, offset=GUARD_ELEMENTS + ) + + out_1d = pto.alloc_tile(shape=[rows, cols], dtype=pto.i32, addr=0) + out_2d = pto.alloc_tile( + shape=list(TEXPANDS_PADDED_SHAPE), + dtype=pto.i32, + valid_shape=[rows, cols], + addr=4096, + ) + + pto.tile.expands(23, out_1d) + pto.tile.expands(23, out_2d) + pto.tile.store(out_1d, out_1d_view) + pto.tile.store(out_2d, out_2d_view) + + +def _equivalence_case( + name, + kernel, + *, + inputs, + expected, + output_dtype, + guard_value, + rtol=0.0, + atol=0.0, +): + def make_case(): + host_inputs = [np.array(value, copy=True) for value in inputs()] + golden = np.array(expected(*host_inputs), dtype=output_dtype, copy=True) + guarded_size = golden.size + 2 * GUARD_ELEMENTS + out_1d = np.full(guarded_size, guard_value, dtype=output_dtype) + out_2d = np.full(guarded_size, guard_value, dtype=output_dtype) + return [*host_inputs, out_1d, out_2d], golden + + def check(device_inputs, golden): + outputs = [ + device_inputs[-2].cpu().numpy(), + device_inputs[-1].cpu().numpy(), + ] + expected_guard = np.full( + GUARD_ELEMENTS, + guard_value, + dtype=output_dtype, + ) + logical_outputs = [] + for output in outputs: + assert_close( + output[:GUARD_ELEMENTS], + expected_guard, + rtol=0.0, + atol=0.0, + ) + assert_close( + output[-GUARD_ELEMENTS:], + expected_guard, + rtol=0.0, + atol=0.0, + ) + logical = output[ + GUARD_ELEMENTS:-GUARD_ELEMENTS + ].reshape(golden.shape) + assert_close(logical, golden, rtol=rtol, atol=atol) + logical_outputs.append(logical) + assert_close( + logical_outputs[0], + logical_outputs[1], + rtol=rtol, + atol=atol, + ) + + return { + "name": name, + "kernel": kernel, + "make_case": make_case, + "check": check, + } + + +def _tabs_inputs(): + rng = np.random.default_rng(SEED + 1) + return [rng.uniform(-20.0, 20.0, size=TABS_SHAPE).astype(np.float32)] + + +def _tadd_inputs(): + rng = np.random.default_rng(SEED + 2) + return [ + rng.integers(-100, 100, size=TADD_SHAPE, dtype=np.int16), + rng.integers(-100, 100, size=TADD_SHAPE, dtype=np.int16), + ] + + +def _tadds_inputs(): + rng = np.random.default_rng(SEED + 3) + return [rng.integers(-50, 50, size=TADDS_SHAPE, dtype=np.int8)] + + +def _tcmps_inputs(): + rng = np.random.default_rng(SEED + 4) + src = rng.integers(0, 10, size=TCMPS_DATA_SHAPE, dtype=np.int8) + return [src] + + +def _tsels_inputs(): + rng = np.random.default_rng(SEED + 5) + predicate = rng.integers( + 0, + 2, + size=TSELS_DATA_SHAPE, + dtype=np.uint8, + ) + mask = np.packbits(predicate, axis=1, bitorder="little").view(np.int8) + src = rng.integers(-20, 20, size=TSELS_DATA_SHAPE, dtype=np.int8) + return [mask, src] + + +def _tcvt_inputs(): + rng = np.random.default_rng(SEED + 6) + src = rng.integers(-1000, 1000, size=TCVT_SHAPE).astype(np.float32) + return [src] + + +CASES = [ + _equivalence_case( + "elementwise_tabs_f32_multirow_tail_1d_2d_equivalence", + elementwise_tabs_1d_2d_equivalence, + inputs=_tabs_inputs, + expected=lambda src: np.abs(src), + output_dtype=np.float32, + guard_value=np.float32(12345.0), + ), + _equivalence_case( + "elementwise_tadd_i16_tail_1d_2d_equivalence", + elementwise_tadd_1d_2d_equivalence, + inputs=_tadd_inputs, + expected=lambda lhs, rhs: (lhs + rhs).astype(np.int16), + output_dtype=np.int16, + guard_value=np.int16(12345), + ), + _equivalence_case( + "elementwise_tadds_i8_tail_1d_2d_equivalence", + elementwise_tadds_1d_2d_equivalence, + inputs=_tadds_inputs, + expected=lambda src: (src.astype(np.int16) + 7).astype(np.int8), + output_dtype=np.int8, + guard_value=np.int8(101), + ), + _equivalence_case( + "elementwise_tcmps_i8_aligned_1d_2d_equivalence", + elementwise_tcmps_1d_2d_equivalence, + inputs=_tcmps_inputs, + expected=lambda src: np.packbits( + src == np.int8(5), + axis=1, + bitorder="little", + ), + output_dtype=np.uint8, + guard_value=np.uint8(0xA5), + ), + _equivalence_case( + "elementwise_tsels_i8_aligned_1d_2d_equivalence", + elementwise_tsels_1d_2d_equivalence, + inputs=_tsels_inputs, + expected=lambda mask, src: np.where( + np.unpackbits( + mask.view(np.uint8), + axis=1, + count=TSELS_DATA_SHAPE[1], + bitorder="little", + ).astype(bool), + src, + np.int8(-3), + ), + output_dtype=np.int8, + guard_value=np.int8(101), + ), + _equivalence_case( + "elementwise_tcvt_f32_i16_tail_1d_2d_equivalence", + elementwise_tcvt_1d_2d_equivalence, + inputs=_tcvt_inputs, + expected=lambda src: src.astype(np.int16), + output_dtype=np.int16, + guard_value=np.int16(12345), + ), + _equivalence_case( + "elementwise_texpands_i32_tail_1d_2d_equivalence", + elementwise_texpands_1d_2d_equivalence, + inputs=lambda: [], + expected=lambda: np.full(TEXPANDS_SHAPE, 23, dtype=np.int32), + output_dtype=np.int32, + guard_value=np.int32(123456789), + ), +] + + +auto_main(globals()) From 0969aa9dc1494c117d83a60931bae13d7b66f3d6 Mon Sep 17 00:00:00 2001 From: Zhendong404 Date: Tue, 11 Aug 2026 19:52:55 +0800 Subject: [PATCH 089/122] fix(ptodsl): pass selection tmp as keyword --- ptodsl/ptodsl/_ops.py | 4 ++-- ptodsl/tests/test_vector_cube_ops.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 65d31d6166..4fb41199da 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -4049,8 +4049,8 @@ def tsel(mask, src0, src1, dst, *, tmp=None): unwrap_surface_value(mask), unwrap_surface_value(src0), unwrap_surface_value(src1), - unwrap_surface_value(resolved_tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(resolved_tmp), ) @@ -4060,9 +4060,9 @@ def tsels(mask, src, scalar, dst, *, tmp=None): _pto.tsels( unwrap_surface_value(mask), unwrap_surface_value(src), - unwrap_surface_value(resolved_tmp), _coerce_tile_scalar_operand(src, scalar, context="tsels"), unwrap_surface_value(dst), + tmp=unwrap_surface_value(resolved_tmp), ) diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index 6f1e0c67b4..51e307bb7f 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -877,23 +877,23 @@ def test_tile_selection_wrappers_use_explicit_tmp_or_synthesize_one(self): patch.object(_ops._pto, "tsel") as tsel_op: _ops.tsel(mask, src0, src1, dst) resolve_tmp.assert_called_once_with(dst, None, context="tsel") - self.assertEqual(tsel_op.call_args.args, (mask, src0, src1, synthesized_tmp, dst)) + tsel_op.assert_called_once_with(mask, src0, src1, dst, tmp=synthesized_tmp) with patch.object(_ops, "_resolve_selection_tmp", side_effect=AssertionError("should not synthesize")), \ patch.object(_ops._pto, "tsel") as tsel_op: _ops.tsel(mask, src0, src1, dst, tmp=tmp) - self.assertEqual(tsel_op.call_args.args, (mask, src0, src1, tmp, dst)) + tsel_op.assert_called_once_with(mask, src0, src1, dst, tmp=tmp) with patch.object(_ops, "_resolve_selection_tmp", return_value=synthesized_tmp) as resolve_tmp, \ patch.object(_ops._pto, "tsels") as tsels_op: _ops.tsels(mask, src, scalar, dst) resolve_tmp.assert_called_once_with(dst, None, context="tsels") - self.assertEqual(tsels_op.call_args.args, (mask, src, synthesized_tmp, coerced_scalar, dst)) + tsels_op.assert_called_once_with(mask, src, coerced_scalar, dst, tmp=synthesized_tmp) with patch.object(_ops, "_resolve_selection_tmp", side_effect=AssertionError("should not synthesize")), \ patch.object(_ops._pto, "tsels") as tsels_op: _ops.tsels(mask, src, scalar, dst, tmp=tmp) - self.assertEqual(tsels_op.call_args.args, (mask, src, tmp, coerced_scalar, dst)) + tsels_op.assert_called_once_with(mask, src, coerced_scalar, dst, tmp=tmp) def test_tile_row_reductions_expose_optional_tmp_and_synthesize_one(self): src = SimpleNamespace(type="src_ty") From 5d9001298fac2f1f406410854d440791df2b7c99 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Mon, 10 Aug 2026 17:42:22 +0800 Subject: [PATCH 090/122] feat(vmi): add carry-chain operations --- docs/isa/vmi-isa/03-eltwise-compute.md | 20 ++ docs/isa/vmi-isa/10-appendices.md | 2 + include/PTO/IR/VMIOps.td | 42 ++++ lib/PTO/IR/VMI.cpp | 39 ++++ lib/PTO/Transforms/VMILayoutAssignment.cpp | 14 ++ lib/PTO/Transforms/VMILayoutPropagation.cpp | 1 + .../Transforms/VMILowerUnifiedToLegacy.cpp | 8 +- .../VMIMaskGranularityAssignment.cpp | 13 ++ lib/PTO/Transforms/VMIToVPTO.cpp | 204 ++++++++++++++++++ ptodsl/ptodsl/_vmi_namespace.py | 58 +++++ ptodsl/tests/test_jit_compile.py | 6 + ptodsl/tests/test_vmi_binary_ops.py | 15 ++ ptodsl/tests/test_vmi_isa_inventory.py | 6 +- .../vmi_new/vmi_carry_verifier_invalid.pto | 38 ++++ test/lit/vmi_new/vmi_to_vpto_carry.pto | 49 +++++ 15 files changed, 509 insertions(+), 6 deletions(-) create mode 100644 test/lit/vmi_new/vmi_carry_verifier_invalid.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_carry.pto diff --git a/docs/isa/vmi-isa/03-eltwise-compute.md b/docs/isa/vmi-isa/03-eltwise-compute.md index 2f6dff4278..cb0ec4e9d1 100644 --- a/docs/isa/vmi-isa/03-eltwise-compute.md +++ b/docs/isa/vmi-isa/03-eltwise-compute.md @@ -68,6 +68,26 @@ : !pto.vmi.vreg<64×f32>, !pto.vmi.vreg<64×f32>, !pto.vmi.mask<64> -> !pto.vmi.vreg<64×f32> ``` +### `pto.vmi.vaddc` / `pto.vmi.vaddcs` + +Carry-chain integer adds are exposed as multi-result VMI operations so the +frontend can preserve the hardware carry instruction instead of expanding the +operation into an add/compare/select sequence. + +```mlir +%sum, %carry = pto.vmi.vaddc %lhs, %rhs, %mask + : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask + -> !pto.vmi.vreg, !pto.vmi.mask +%next, %carry2 = pto.vmi.vaddcs %lhs, %rhs, %carry, %mask + : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask, !pto.vmi.mask + -> !pto.vmi.vreg, !pto.vmi.mask +``` + +Both operations require matching 32-bit integer data values. The execution +mask, carry-in (for `vaddcs`), and carry-out use the same logical lane count, +layout, and `b32` physical mask granularity as the data ports. They lower +one-to-N to `pto.vaddc` and `pto.vaddcs` respectively. + ### `pto.vmi.vdiv` - **semantics:** Elementwise floating-point divide. diff --git a/docs/isa/vmi-isa/10-appendices.md b/docs/isa/vmi-isa/10-appendices.md index a0586589b7..597bcc50ff 100644 --- a/docs/isa/vmi-isa/10-appendices.md +++ b/docs/isa/vmi-isa/10-appendices.md @@ -59,6 +59,8 @@ | 51 | `pto.vmi.create_group_mask` | 8: Predicate | gen | Grouped predicate mask | | 52 | `pto.vmi.vintlv` | 9: Rearrange | A | Interleave two vectors | | 53 | `pto.vmi.vdintlv` | 9: Rearrange | A | Deinterleave two vectors | +| 54 | `pto.vmi.vaddc` | 3: Eltwise | A | 32-bit integer add with per-lane carry output | +| 55 | `pto.vmi.vaddcs` | 3: Eltwise | A | 32-bit integer add with carry input and output | --- diff --git a/include/PTO/IR/VMIOps.td b/include/PTO/IR/VMIOps.td index e94c49870d..1f21c55b67 100644 --- a/include/PTO/IR/VMIOps.td +++ b/include/PTO/IR/VMIOps.td @@ -953,6 +953,48 @@ def VMIVaddOp : VMI_Op<"vadd", [Pure]> { let assemblyFormat = "$lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; } +def VMIVaddcOp : VMI_Op<"vaddc", [Pure]> { + let summary = "VMI integer add with per-lane carry output"; + let description = [{ + Adds two integer vectors and returns both the truncated result and the + per-lane carry predicate. The operation is restricted to the 32-bit + integer forms supported by the underlying VPTO carry instruction. + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$lhs, + VMI_VRegTypeConstraint:$rhs, + VMI_MaskTypeConstraint:$mask + ); + let results = (outs + VMI_VRegTypeConstraint:$result, + VMI_MaskTypeConstraint:$carry + ); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs `,` $mask attr-dict `:` type($lhs) `,` type($rhs) `,` type($mask) `->` type($result) `,` type($carry)"; +} + +def VMIVaddcsOp : VMI_Op<"vaddcs", [Pure]> { + let summary = "VMI integer add with carry-in and carry-out"; + let description = [{ + Adds two integer vectors and a per-lane carry-in predicate, returning both + the truncated result and the per-lane carry-out predicate. The operation + is restricted to the 32-bit integer forms supported by the underlying VPTO + carry instruction. + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$lhs, + VMI_VRegTypeConstraint:$rhs, + VMI_MaskTypeConstraint:$carry_in, + VMI_MaskTypeConstraint:$mask + ); + let results = (outs + VMI_VRegTypeConstraint:$result, + VMI_MaskTypeConstraint:$carry + ); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs `,` $carry_in `,` $mask attr-dict `:` type($lhs) `,` type($rhs) `,` type($carry_in) `,` type($mask) `->` type($result) `,` type($carry)"; +} + def VMIVsubOp : VMI_Op<"vsub", [Pure]> { let summary = "VMI elementwise subtract (unified fp/int)"; let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs, diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp index c59a8946b7..c10315d918 100644 --- a/lib/PTO/IR/VMI.cpp +++ b/lib/PTO/IR/VMI.cpp @@ -3543,6 +3543,45 @@ LogicalResult VMIVpreluOp::verify() { return success(); } +template +static LogicalResult verifyVMIAddCarryOp(CarryOp op, VMIMaskType carryInType, + bool hasCarryIn) { + auto lhsType = cast(op.getLhs().getType()); + auto rhsType = cast(op.getRhs().getType()); + auto resultType = cast(op.getResult().getType()); + auto maskType = cast(op.getMask().getType()); + auto carryType = cast(op.getCarry().getType()); + + auto integerType = dyn_cast(lhsType.getElementType()); + if (!integerType || integerType.getWidth() != 32) + return op.emitOpError("requires 32-bit integer vector element types"); + + if (failed(verifyAllSameVRegShapeAndLayout( + op.getOperation(), {lhsType, rhsType, resultType}, + /*requireSameElement=*/true))) + return failure(); + if (failed(verifyMaskMatchesData(op.getOperation(), maskType, lhsType)) || + failed(verifyMaskMatchesData(op.getOperation(), carryType, lhsType))) + return failure(); + if (hasCarryIn && + failed(verifyMaskMatchesData(op.getOperation(), carryInType, lhsType))) + return failure(); + + SmallVector masks{maskType, carryType}; + if (hasCarryIn) + masks.push_back(carryInType); + return verifyAllSameMaskShapeLayoutAndGranularity(op.getOperation(), masks); +} + +LogicalResult VMIVaddcOp::verify() { + return verifyVMIAddCarryOp(*this, VMIMaskType{}, /*hasCarryIn=*/false); +} + +LogicalResult VMIVaddcsOp::verify() { + return verifyVMIAddCarryOp(*this, cast(getCarryIn().getType()), + /*hasCarryIn=*/true); +} + LogicalResult VMIVmullOp::verify() { auto aType = cast(getA().getType()); auto bType = cast(getB().getType()); diff --git a/lib/PTO/Transforms/VMILayoutAssignment.cpp b/lib/PTO/Transforms/VMILayoutAssignment.cpp index 89c58f56fd..49555602fc 100644 --- a/lib/PTO/Transforms/VMILayoutAssignment.cpp +++ b/lib/PTO/Transforms/VMILayoutAssignment.cpp @@ -655,6 +655,20 @@ struct LayoutSolver { return WalkResult::interrupt(); return WalkResult::advance(); } + if (auto addc = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(addc.getLhsMutable(), + addc.getRhsMutable(), + addc.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto addcs = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(addcs.getLhsMutable(), + addcs.getRhsMutable(), + addcs.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } if (auto fma = dyn_cast(op)) { if (failed(unite(fma.getLhs(), fma.getRhs(), op)) || failed(unite(fma.getLhs(), fma.getAcc(), op)) || diff --git a/lib/PTO/Transforms/VMILayoutPropagation.cpp b/lib/PTO/Transforms/VMILayoutPropagation.cpp index 0819d42438..d3004656f9 100644 --- a/lib/PTO/Transforms/VMILayoutPropagation.cpp +++ b/lib/PTO/Transforms/VMILayoutPropagation.cpp @@ -170,6 +170,7 @@ class VMILayoutMaterializationTransfer final { static bool isSameLayoutOp(Operation *op) { return isa( op)) { op->emitRemark("VMI unified op has no legacy equivalent — " diff --git a/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp b/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp index 72c6120cdf..3527c61c96 100644 --- a/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp +++ b/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp @@ -272,6 +272,19 @@ struct MaskGranularitySolver { return WalkResult::interrupt(); return WalkResult::advance(); } + if (auto addc = dyn_cast(op)) { + if (failed(requestMaskUse(addc.getMaskMutable(), "b32", op)) || + failed(requestMask(addc.getCarry(), "b32", op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto addcs = dyn_cast(op)) { + if (failed(requestMaskUse(addcs.getCarryInMutable(), "b32", op)) || + failed(requestMaskUse(addcs.getMaskMutable(), "b32", op)) || + failed(requestMask(addcs.getCarry(), "b32", op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } if (auto activePrefix = dyn_cast(op)) { auto resultType = cast(activePrefix.getResult().getType()); if (failed(requestMaskUse( diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index d0a2288dc6..043329ce59 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -8610,6 +8610,123 @@ struct OneToNVMIVecScalarOpPattern : OpConversionPattern { } }; +struct OneToNVMIVaddcOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIVaddcOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange lhsParts = adaptor.getLhs(); + ValueRange rhsParts = adaptor.getRhs(); + ValueRange maskParts = adaptor.getMask(); + FailureOr> maybeResultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + FailureOr> maybeCarryTypes = + getConvertedResultTypes(op, 1, *this->getTypeConverter()); + if (failed(maybeResultTypes) || failed(maybeCarryTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybeResultTypes); + SmallVector carryTypes = std::move(*maybeCarryTypes); + + if (lhsParts.empty() || rhsParts.size() != lhsParts.size() || + maskParts.size() != lhsParts.size() || + resultTypes.size() != lhsParts.size() || + carryTypes.size() != lhsParts.size()) + return rewriter.notifyMatchFailure(op, + "vaddc physical arity mismatch"); + + SmallVector results; + SmallVector carries; + results.reserve(lhsParts.size()); + carries.reserve(lhsParts.size()); + for (auto [lhs, rhs, mask, resultType, carryType] : + llvm::zip_equal(lhsParts, rhsParts, maskParts, resultTypes, + carryTypes)) { + auto dataType = dyn_cast(resultType); + auto integerType = dataType + ? dyn_cast(dataType.getElementType()) + : IntegerType(); + if (!dataType || !integerType || integerType.getWidth() != 32 || + !isa(mask.getType()) || + !isa(carryType) || !cast(carryType).isB32() || + lhs.getType() != resultType || rhs.getType() != resultType) + return rewriter.notifyMatchFailure( + op, "vaddc requires matching 32-bit data and b32 mask parts"); + + auto addc = rewriter.create(op.getLoc(), resultType, carryType, + lhs, rhs, mask); + results.push_back(addc.getResult()); + carries.push_back(addc.getCarry()); + } + + results.append(carries); + replaceOpWithFlatConvertedValues(rewriter, op, results, + *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIVaddcsOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIVaddcsOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange lhsParts = adaptor.getLhs(); + ValueRange rhsParts = adaptor.getRhs(); + ValueRange carryInParts = adaptor.getCarryIn(); + ValueRange maskParts = adaptor.getMask(); + FailureOr> maybeResultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + FailureOr> maybeCarryTypes = + getConvertedResultTypes(op, 1, *this->getTypeConverter()); + if (failed(maybeResultTypes) || failed(maybeCarryTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybeResultTypes); + SmallVector carryTypes = std::move(*maybeCarryTypes); + + if (lhsParts.empty() || rhsParts.size() != lhsParts.size() || + carryInParts.size() != lhsParts.size() || + maskParts.size() != lhsParts.size() || + resultTypes.size() != lhsParts.size() || + carryTypes.size() != lhsParts.size()) + return rewriter.notifyMatchFailure(op, + "vaddcs physical arity mismatch"); + + SmallVector results; + SmallVector carries; + results.reserve(lhsParts.size()); + carries.reserve(lhsParts.size()); + for (auto [lhs, rhs, carryIn, mask, resultType, carryType] : + llvm::zip_equal(lhsParts, rhsParts, carryInParts, maskParts, + resultTypes, carryTypes)) { + auto dataType = dyn_cast(resultType); + auto integerType = dataType + ? dyn_cast(dataType.getElementType()) + : IntegerType(); + if (!dataType || !integerType || integerType.getWidth() != 32 || + !isa(carryIn.getType()) || + !isa(mask.getType()) || !isa(carryType) || + !cast(carryIn.getType()).isB32() || + !cast(mask.getType()).isB32() || + !cast(carryType).isB32() || + lhs.getType() != resultType || rhs.getType() != resultType) + return rewriter.notifyMatchFailure( + op, "vaddcs requires matching 32-bit data and b32 mask parts"); + + auto addcs = rewriter.create(op.getLoc(), resultType, carryType, + lhs, rhs, carryIn, mask); + results.push_back(addcs.getResult()); + carries.push_back(addcs.getCarry()); + } + + results.append(carries); + replaceOpWithFlatConvertedValues(rewriter, op, results, + *this->getTypeConverter()); + return success(); + } +}; + struct OneToNVMIVmullOpPattern : OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -12164,6 +12281,7 @@ void populateVMIConversionPatterns( OneToNVMIMaskedStoreOpPattern, OneToNVMIStrideStoreOpPattern, OneToNVMIScatterOpPattern, OneToNVMIBinaryOpPattern, OneToNVMIBinaryOpPattern, + OneToNVMIVaddcOpPattern, OneToNVMIVaddcsOpPattern, OneToNVMIBinaryOpPattern, OneToNVMIBinaryOpPattern, OneToNVMIBinaryOpPattern, @@ -12890,6 +13008,72 @@ LogicalResult checkSupportedVmullShape(VMIVmullOp op, return success(); } +static LogicalResult +checkSupportedVMIAddCarryPorts(VMIVRegType lhsType, VMIVRegType rhsType, + VMIVRegType resultType, + ArrayRef maskTypes, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto integerType = dyn_cast(lhsType.getElementType()); + if (!integerType || integerType.getWidth() != 32) + return fail("requires 32-bit integer data elements"); + if (lhsType != rhsType || lhsType != resultType) + return fail("requires matching lhs, rhs, and result VMI types"); + if (!lhsType.getLayoutAttr()) + return fail("requires assigned data layout"); + if (failed(checkSupportedMaskableVReg(lhsType))) + return fail("requires computable physical data parts"); + + FailureOr dataArity = getVMIPhysicalArity(lhsType); + if (failed(dataArity) || *dataArity < 1) + return fail("requires non-empty physical data parts"); + for (VMIMaskType maskType : maskTypes) { + if (maskType.getLayoutAttr() != lhsType.getLayoutAttr()) + return fail("requires all data and mask ports to share one layout"); + if (maskType.getGranularity() != "b32") + return fail("requires b32 mask granularity"); + FailureOr maskArity = getVMIPhysicalArity(maskType); + if (failed(maskArity) || *maskArity != *dataArity) + return fail("requires matching physical arity on data and mask ports"); + FailureOr physicalGranularity = + getVMIMaskPhysicalGranularity(maskType); + if (failed(physicalGranularity) || *physicalGranularity != "b32") + return fail("requires physical b32 mask parts"); + } + FailureOr lanesPerPart = getDataLanesPerPart(lhsType.getElementType()); + if (failed(lanesPerPart) || *lanesPerPart != 64) + return fail("requires 64-lane 32-bit data parts"); + return success(); +} + +LogicalResult checkSupportedVMIAddcShape(VMIVaddcOp op, + std::string *reason = nullptr) { + return checkSupportedVMIAddCarryPorts( + cast(op.getLhs().getType()), + cast(op.getRhs().getType()), + cast(op.getResult().getType()), + {cast(op.getMask().getType()), + cast(op.getCarry().getType())}, + reason); +} + +LogicalResult checkSupportedVMIAddcsShape(VMIVaddcsOp op, + std::string *reason = nullptr) { + return checkSupportedVMIAddCarryPorts( + cast(op.getLhs().getType()), + cast(op.getRhs().getType()), + cast(op.getResult().getType()), + {cast(op.getCarryIn().getType()), + cast(op.getMask().getType()), + cast(op.getCarry().getType())}, + reason); +} + LogicalResult checkSupportedFmaShape(VMIFmaOp op, std::string *reason = nullptr) { auto fail = [&](const Twine &message) -> LogicalResult { @@ -13307,6 +13491,26 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto muli = dyn_cast(op)) return emitMaskableUnsupported( op, "pto.vmi.muli", cast(muli.getResult().getType())); + if (auto addc = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedVMIAddcShape(addc, &reason))) + return WalkResult::advance(); + addc.emitError() << kVMIDiagUnsupportedPrefix + << "pto.vmi.vaddc requires matching 32-bit data and " + "b32 mask parts (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto addcs = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedVMIAddcsShape(addcs, &reason))) + return WalkResult::advance(); + addcs.emitError() << kVMIDiagUnsupportedPrefix + << "pto.vmi.vaddcs requires matching 32-bit data and " + "b32 mask parts (" + << reason << ")"; + return WalkResult::interrupt(); + } auto verifyVecScalar = [&](auto vecScalar, StringRef opName) -> WalkResult { if (vecScalar.getPmode().has_value() && *vecScalar.getPmode() == "merge") { diff --git a/ptodsl/ptodsl/_vmi_namespace.py b/ptodsl/ptodsl/_vmi_namespace.py index 01093ca10c..acbdef830b 100644 --- a/ptodsl/ptodsl/_vmi_namespace.py +++ b/ptodsl/ptodsl/_vmi_namespace.py @@ -341,6 +341,25 @@ def _derive_vmull_result_types(a, b, *, context: str): return lhs_type, rhs_type +def _derive_add_carry_result_types(lhs, rhs, mask, *, carry_in=None, context: str): + lhs_type = _as_vmi_vreg_type(_type_of(lhs), context=context) + rhs_type = _as_vmi_vreg_type(_type_of(rhs), context=context) + if lhs_type != rhs_type: + raise TypeError(f"{context} requires lhs and rhs to have identical VMI vreg types") + element_type = lhs_type.element_type + if not IntegerType.isinstance(element_type) or IntegerType(element_type).width != 32: + raise TypeError(f"{context} requires 32-bit integer vectors") + + mask_type = _as_vmi_mask_type(_type_of(mask), context=context) + if _vmi_mask_element_count(mask_type, context=context) != lhs_type.element_count: + raise TypeError(f"{context} requires the mask lane count to match the data vectors") + if carry_in is not None: + carry_in_type = _as_vmi_mask_type(_type_of(carry_in), context=context) + if carry_in_type != mask_type: + raise TypeError(f"{context} requires carry_in and mask to have identical VMI mask types") + return lhs_type, mask_type + + def _derive_hist_result_type(acc, *, context: str): """acc must be 16-bit unsigned or signless integer; result is always ui16.""" acc_type = _as_vmi_vreg_type(_type_of(acc), context=context) @@ -761,6 +780,45 @@ def vadd(lhs, rhs, mask=None, **kw): """Emit VMI vector addition, selecting vector or scalar form by type.""" return _emit_binary_or_vec_scalar("vadd", "vadds", lhs, rhs, mask, commutative=True, **kw) + @staticmethod + def vaddc(lhs, rhs, mask, *, loc=None, ip=None): + """Emit a 32-bit integer add with per-lane carry output.""" + context = "pto.vmi.vaddc(...)" + mask_value = _required_mask(mask, context=context) + result_type, carry_type = _derive_add_carry_result_types( + lhs, rhs, mask_value, context=context + ) + return _call_value( + "vaddc", + result_type, + carry_type, + _raw(lhs), + _raw(rhs), + mask_value, + loc=loc, + ip=ip, + ) + + @staticmethod + def vaddcs(lhs, rhs, carry_in, mask, *, loc=None, ip=None): + """Emit a 32-bit integer add with carry input and carry output.""" + context = "pto.vmi.vaddcs(...)" + mask_value = _required_mask(mask, context=context) + result_type, carry_type = _derive_add_carry_result_types( + lhs, rhs, mask_value, carry_in=carry_in, context=context + ) + return _call_value( + "vaddcs", + result_type, + carry_type, + _raw(lhs), + _raw(rhs), + _raw(carry_in), + mask_value, + loc=loc, + ip=ip, + ) + vsub = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vsub", lhs, rhs, mask, **kw)) @staticmethod diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 20f8acb846..c2d6bdc708 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -2679,6 +2679,8 @@ def vmi_wrapper_dispatch_probe(): int_rhs = pto.vmi.vload(int_other_ptr, offset, size=64) hist_mask = pto.vmi.create_mask(pto.const(256, dtype=pto.index), size=256) added = pto.vmi.vadd(lhs, rhs, mask) + carry_sum, carry = pto.vmi.vaddc(int_lhs, int_rhs, mask) + carry_next, carry_out = pto.vmi.vaddcs(carry_sum, int_rhs, carry, mask) subtracted = pto.vmi.vsub(lhs, rhs, mask) multiplied = pto.vmi.vmul(lhs, rhs, mask) divided = pto.vmi.vdiv(lhs, rhs, mask) @@ -2742,6 +2744,8 @@ def vmi_wrapper_dispatch_probe(): pto.vmi.vsstb(hi, dst_ptr, offset, pto.i16(8), mask) _ = group_mask + _ = carry_next + _ = carry_out _ = total _ = explicit_total _ = peak @@ -6594,6 +6598,8 @@ def _enter_inline_simt_with_resource_attr(): "pto.vmi.vsstb", "pto.vmi.vci", "pto.vmi.vadd", + "pto.vmi.vaddc", + "pto.vmi.vaddcs", "pto.vmi.vsub", "pto.vmi.vmul", "pto.vmi.vdiv", diff --git a/ptodsl/tests/test_vmi_binary_ops.py b/ptodsl/tests/test_vmi_binary_ops.py index 47ccdf889b..a7a9ee1ae0 100644 --- a/ptodsl/tests/test_vmi_binary_ops.py +++ b/ptodsl/tests/test_vmi_binary_ops.py @@ -79,6 +79,17 @@ def vmi_binary_add_compatibility_probe(): _ = pto.vmi.vadds(source, 1.0, mask) +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_add_carry_probe(): + lhs_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.ui32) + rhs_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.ui32) + mask = pto.vmi.create_mask(64, size=64) + lhs = pto.vmi.vload(lhs_tile.as_ptr(), 0, size=64) + rhs = pto.vmi.vload(rhs_tile.as_ptr(), 0, size=64) + sum_value, carry = pto.vmi.vaddc(lhs, rhs, mask) + _, _ = pto.vmi.vaddcs(sum_value, rhs, carry, mask) + + def expect(condition: bool, message: str) -> None: if not condition: raise AssertionError(message) @@ -101,6 +112,10 @@ def main() -> None: "vmi.vadd(vector, scalar, mask) should emit pto.vmi.vadds", ) + carry_text = vmi_add_carry_probe.compile().mlir_text() + expect("pto.vmi.vaddc" in carry_text, "vmi.vaddc should emit pto.vmi.vaddc") + expect("pto.vmi.vaddcs" in carry_text, "vmi.vaddcs should emit pto.vmi.vaddcs") + vector_scalar_text = vmi_binary_vector_scalar_probe.compile().mlir_text() for op_name in ("vmuls", "vmaxs", "vmins", "vshls", "vshrs"): expect( diff --git a/ptodsl/tests/test_vmi_isa_inventory.py b/ptodsl/tests/test_vmi_isa_inventory.py index 49d7b43ada..4734da87cd 100644 --- a/ptodsl/tests/test_vmi_isa_inventory.py +++ b/ptodsl/tests/test_vmi_isa_inventory.py @@ -24,7 +24,7 @@ BACKEND_CAPABILITY = { name: "a5-vpto" for name in ( - "vload vstore vsstb vci vadd vsub vmul vdiv vmax vmin vabs vneg " + "vload vstore vsstb vci vadd vaddc vaddcs vsub vmul vdiv vmax vmin vabs vneg " "vrelu vexp vln vsqrt vand vor vxor vnot vshl vshr vadds vmuls " "vmaxs vmins vshls vshrs vcmp vcmps vsel vselr vbrc vcadd vcmax " "vcmin vcvt vinterpret_cast vexpdif vaxpy vlrelu vprelu vmull " @@ -48,9 +48,9 @@ def _indexed_ops(): def main() -> None: indexed = _indexed_ops() - assert [number for number, _ in indexed] == list(range(1, 54)) + assert [number for number, _ in indexed] == list(range(1, 56)) names = [name for _, name in indexed] - assert len(names) == len(set(names)) == 53 + assert len(names) == len(set(names)) == 55 assert set(BACKEND_CAPABILITY) == set(names) assert set(PTODSL_ALIASES) <= set(names) diff --git a/test/lit/vmi_new/vmi_carry_verifier_invalid.pto b/test/lit/vmi_new/vmi_carry_verifier_invalid.pto new file mode 100644 index 0000000000..e4d7820ff1 --- /dev/null +++ b/test/lit/vmi_new/vmi_carry_verifier_invalid.pto @@ -0,0 +1,38 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not pto-test-opt %s -split-input-file 2>&1 | FileCheck %s + +module { + func.func @vaddc_f16_invalid( + %lhs: !pto.vmi.vreg<64xf16>, %rhs: !pto.vmi.vreg<64xf16>, + %mask: !pto.vmi.mask<64xpred>) { + %sum, %carry = pto.vmi.vaddc %lhs, %rhs, %mask + : !pto.vmi.vreg<64xf16>, !pto.vmi.vreg<64xf16>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xf16>, !pto.vmi.mask<64xpred> + return + } +} + +// CHECK: 'pto.vmi.vaddc' op requires 32-bit integer vector element types + +// ----- + +module { + func.func @vaddcs_carry_lane_mismatch( + %lhs: !pto.vmi.vreg<64xui32>, %rhs: !pto.vmi.vreg<64xui32>, + %carry_in: !pto.vmi.mask<32xpred>, %mask: !pto.vmi.mask<64xpred>) { + %sum, %carry = pto.vmi.vaddcs %lhs, %rhs, %carry_in, %mask + : !pto.vmi.vreg<64xui32>, !pto.vmi.vreg<64xui32>, + !pto.vmi.mask<32xpred>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xui32>, !pto.vmi.mask<64xpred> + return + } +} + +// CHECK: 'pto.vmi.vaddcs' op requires mask logical lane count to match data lane count diff --git a/test/lit/vmi_new/vmi_to_vpto_carry.pto b/test/lit/vmi_new/vmi_to_vpto_carry.pto new file mode 100644 index 0000000000..803f5b105e --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_carry.pto @@ -0,0 +1,49 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +module { + func.func @vmi_to_vpto_vaddc( + %lhs: !pto.vmi.vreg<128xui32>, + %rhs: !pto.vmi.vreg<128xui32>, + %mask: !pto.vmi.mask<128xpred>) + -> (!pto.vmi.vreg<128xui32>, !pto.vmi.mask<128xpred>) { + %sum, %carry = pto.vmi.vaddc %lhs, %rhs, %mask + : !pto.vmi.vreg<128xui32>, !pto.vmi.vreg<128xui32>, + !pto.vmi.mask<128xpred> + -> !pto.vmi.vreg<128xui32>, !pto.vmi.mask<128xpred> + return %sum, %carry + : !pto.vmi.vreg<128xui32>, !pto.vmi.mask<128xpred> + } + + func.func @vmi_to_vpto_vaddcs( + %lhs: !pto.vmi.vreg<128xui32>, + %rhs: !pto.vmi.vreg<128xui32>, + %carry_in: !pto.vmi.mask<128xpred>, + %mask: !pto.vmi.mask<128xpred>) + -> (!pto.vmi.vreg<128xui32>, !pto.vmi.mask<128xpred>) { + %sum, %carry = pto.vmi.vaddcs %lhs, %rhs, %carry_in, %mask + : !pto.vmi.vreg<128xui32>, !pto.vmi.vreg<128xui32>, + !pto.vmi.mask<128xpred>, !pto.vmi.mask<128xpred> + -> !pto.vmi.vreg<128xui32>, !pto.vmi.mask<128xpred> + return %sum, %carry + : !pto.vmi.vreg<128xui32>, !pto.vmi.mask<128xpred> + } +} + +// CHECK-LABEL: func.func @vmi_to_vpto_vaddc( +// CHECK: %[[SUM0:.*]], %[[CARRY0:.*]] = pto.vaddc +// CHECK: %[[SUM1:.*]], %[[CARRY1:.*]] = pto.vaddc +// CHECK: return %[[SUM0]], %[[SUM1]], %[[CARRY0]], %[[CARRY1]] +// CHECK-LABEL: func.func @vmi_to_vpto_vaddcs( +// CHECK: %[[NEXT0:.*]], %[[NEXTCARRY0:.*]] = pto.vaddcs +// CHECK: %[[NEXT1:.*]], %[[NEXTCARRY1:.*]] = pto.vaddcs +// CHECK: return %[[NEXT0]], %[[NEXT1]], %[[NEXTCARRY0]], %[[NEXTCARRY1]] +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. From 7fba4ceb1cbb75d6c552268812f20df49e605724 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Mon, 10 Aug 2026 19:44:01 +0800 Subject: [PATCH 091/122] fix(vmi): centralize carry layout support --- include/PTO/Transforms/VMILayoutSupport.h | 35 +++ lib/PTO/Transforms/VMILayoutAssignment.cpp | 19 +- lib/PTO/Transforms/VMILayoutPropagation.cpp | 60 ++++- lib/PTO/Transforms/VMILayoutSupport.cpp | 225 ++++++++++++++++++ lib/PTO/Transforms/VMIToVPTO.cpp | 60 +---- .../vmi_carry_layout_support_invalid.pto | 55 +++++ test/lit/vmi_new/vmi_to_vpto_carry.pto | 21 ++ 7 files changed, 414 insertions(+), 61 deletions(-) create mode 100644 test/lit/vmi_new/vmi_carry_layout_support_invalid.pto diff --git a/include/PTO/Transforms/VMILayoutSupport.h b/include/PTO/Transforms/VMILayoutSupport.h index 179640e9c9..bf9fb5d085 100644 --- a/include/PTO/Transforms/VMILayoutSupport.h +++ b/include/PTO/Transforms/VMILayoutSupport.h @@ -104,6 +104,11 @@ struct VMIInterleaveLayoutFact { int64_t lanesPerPart = 0; }; +struct VMIAddCarryLayoutFact { + VMILayoutAttr layout; + int64_t physicalArity = 0; +}; + struct VMIBitcastLayoutFact { VMILayoutAttr sourceLayout; VMILayoutAttr resultLayout; @@ -316,6 +321,30 @@ class VMILayoutSupport { VMIVRegType lowType, VMIVRegType highType, std::string *reason = nullptr) const; + FailureOr + getPreferredVaddcLayoutFact(VMIVaddcOp op, + std::string *reason = nullptr) const; + + FailureOr + getPreferredVaddcsLayoutFact(VMIVaddcsOp op, + std::string *reason = nullptr) const; + + FailureOr> + getVaddcLayoutFactsForLayout(VMIVaddcOp op, VMILayoutAttr layout, + std::string *reason = nullptr) const; + + FailureOr> + getVaddcsLayoutFactsForLayout(VMIVaddcsOp op, VMILayoutAttr layout, + std::string *reason = nullptr) const; + + FailureOr + getVaddcLayoutFact(VMIVaddcOp op, + std::string *reason = nullptr) const; + + FailureOr + getVaddcsLayoutFact(VMIVaddcsOp op, + std::string *reason = nullptr) const; + FailureOr getGroupSlotLoadLayoutFact(VMIVRegType resultType, int64_t numGroups, std::string *reason = nullptr) const; @@ -410,6 +439,12 @@ class VMILayoutSupport { LogicalResult getVselrSupport(VMIVselrOp op, std::string *reason = nullptr) const; + LogicalResult getVaddcSupport(VMIVaddcOp op, + std::string *reason = nullptr) const; + + LogicalResult getVaddcsSupport(VMIVaddcsOp op, + std::string *reason = nullptr) const; + LogicalResult getGroupReduceAddFSupport(VMIGroupReduceAddFOp op, std::string *reason = nullptr) const; diff --git a/lib/PTO/Transforms/VMILayoutAssignment.cpp b/lib/PTO/Transforms/VMILayoutAssignment.cpp index 49555602fc..7710c154ce 100644 --- a/lib/PTO/Transforms/VMILayoutAssignment.cpp +++ b/lib/PTO/Transforms/VMILayoutAssignment.cpp @@ -656,16 +656,31 @@ struct LayoutSolver { return WalkResult::advance(); } if (auto addc = dyn_cast(op)) { + VMILayoutSupport supports; + std::string reason; + if (failed(supports.getPreferredVaddcLayoutFact(addc, &reason))) { + addc.emitError() << kVMIDiagLayoutContractPrefix << reason; + return WalkResult::interrupt(); + } if (failed(constrainElementwiseBinary(addc.getLhsMutable(), addc.getRhsMutable(), - addc.getResult(), op))) + addc.getResult(), op)) || + failed(uniteMask(addc.getMask(), addc.getCarry(), op))) return WalkResult::interrupt(); return WalkResult::advance(); } if (auto addcs = dyn_cast(op)) { + VMILayoutSupport supports; + std::string reason; + if (failed(supports.getPreferredVaddcsLayoutFact(addcs, &reason))) { + addcs.emitError() << kVMIDiagLayoutContractPrefix << reason; + return WalkResult::interrupt(); + } if (failed(constrainElementwiseBinary(addcs.getLhsMutable(), addcs.getRhsMutable(), - addcs.getResult(), op))) + addcs.getResult(), op)) || + failed(uniteMask(addcs.getCarryIn(), addcs.getMask(), op)) || + failed(uniteMask(addcs.getMask(), addcs.getCarry(), op))) return WalkResult::interrupt(); return WalkResult::advance(); } diff --git a/lib/PTO/Transforms/VMILayoutPropagation.cpp b/lib/PTO/Transforms/VMILayoutPropagation.cpp index d3004656f9..78d42b7e5f 100644 --- a/lib/PTO/Transforms/VMILayoutPropagation.cpp +++ b/lib/PTO/Transforms/VMILayoutPropagation.cpp @@ -170,7 +170,6 @@ class VMILayoutMaterializationTransfer final { static bool isSameLayoutOp(Operation *op) { return isa> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + if (auto addc = dyn_cast(op)) + return queryVaddc(addc, changedLayout); + if (auto addcs = dyn_cast(op)) + return queryVaddcs(addcs, changedLayout); + return failure(); + } + +private: + FailureOr> + queryVaddc(VMIVaddcOp op, VMILayoutAttr changedLayout) const { + VMILayoutSupport supports; + FailureOr> facts = + supports.getVaddcLayoutFactsForLayout(op, changedLayout); + if (failed(facts) || facts->empty()) + return failure(); + + SmallVector relations; + for (const VMIAddCarryLayoutFact &fact : *facts) { + relations.push_back(makeRelation(SmallVector{ + operandFact(op.getLhsMutable(), fact.layout), + operandFact(op.getRhsMutable(), fact.layout), + operandFact(op.getMaskMutable(), fact.layout), + valueFact(op.getResult(), fact.layout), + valueFact(op.getCarry(), fact.layout)})); + } + return relations; + } + + FailureOr> + queryVaddcs(VMIVaddcsOp op, VMILayoutAttr changedLayout) const { + VMILayoutSupport supports; + FailureOr> facts = + supports.getVaddcsLayoutFactsForLayout(op, changedLayout); + if (failed(facts) || facts->empty()) + return failure(); + + SmallVector relations; + for (const VMIAddCarryLayoutFact &fact : *facts) { + relations.push_back(makeRelation(SmallVector{ + operandFact(op.getLhsMutable(), fact.layout), + operandFact(op.getRhsMutable(), fact.layout), + operandFact(op.getCarryInMutable(), fact.layout), + operandFact(op.getMaskMutable(), fact.layout), + valueFact(op.getResult(), fact.layout), + valueFact(op.getCarry(), fact.layout)})); + } + return relations; + } +}; + class VMIInterleaveTransfer final : public VMILayoutTransfer { public: FailureOr> @@ -834,6 +889,7 @@ const VMILayoutTransfer *getTransfer(Operation *op) { static VMIGroupSlotLoadTransfer groupSlotLoadTransfer; static VMIGroupBroadcastLoadTransfer groupBroadcastLoadTransfer; static VMIGroupBroadcastTransfer groupBroadcastTransfer; + static VMIAddCarryTransfer addCarryTransfer; static VMIInterleaveTransfer interleaveTransfer; static VMIGatherTransfer gatherTransfer; static VMIStoreTransfer storeTransfer; @@ -862,6 +918,8 @@ const VMILayoutTransfer *getTransfer(Operation *op) { return &groupBroadcastLoadTransfer; if (isa(op)) return &groupBroadcastTransfer; + if (isa(op)) + return &addCarryTransfer; if (isa(op)) return &interleaveTransfer; if (isa(op)) diff --git a/lib/PTO/Transforms/VMILayoutSupport.cpp b/lib/PTO/Transforms/VMILayoutSupport.cpp index 4ec44e2e7f..0b6372f56f 100644 --- a/lib/PTO/Transforms/VMILayoutSupport.cpp +++ b/lib/PTO/Transforms/VMILayoutSupport.cpp @@ -451,6 +451,12 @@ struct InterleaveLayoutPattern { LayoutPattern highLayout; }; +struct AddCarryLayoutPattern { + ElementBitsPattern dataBits; + MaskGranularityPattern maskGranularity; + LayoutPattern layout; +}; + static constexpr PreferredCastLayoutPattern kPreferredCastLayoutPatterns[] = { // Exact rows override the default legal relation for small shapes where the // compact lane-stride form is the natural cast layout. @@ -604,6 +610,13 @@ static constexpr InterleaveLayoutPattern kVintlvLayoutPatterns[] = { {bits<8>(), chunk<1>(), 1, ls(4), ls(4), ls(4), ls(4), ls(4)}, }; +static constexpr AddCarryLayoutPattern kAddCarryLayoutPatterns[] = { + {bits<32>(), mb32(), c()}, {bits<32>(), mb32(), d(2)}, + {bits<32>(), mb32(), d(4)}, {bits<32>(), mb32(), bd(2)}, + {bits<32>(), mb32(), bd(4)}, {bits<32>(), mb32(), gs(1)}, + {bits<32>(), mb32(), gs(8)}, +}; + struct DenseMemoryLayoutPattern { ElementBitsPattern elementBits; LayoutPattern layout; @@ -1122,6 +1135,119 @@ static VMIInterleaveLayoutFact materializeInterleaveLayoutFact( return fact; } +static VMIAddCarryLayoutFact materializeAddCarryLayoutFact( + MLIRContext *ctx, const AddCarryLayoutPattern &pattern, int64_t numGroups, + int64_t physicalArity) { + return VMIAddCarryLayoutFact{ + materializeLayoutPattern(ctx, pattern.layout, numGroups), physicalArity}; +} + +static FailureOr> +getAddCarryLayoutFactsForLayoutImpl( + VMIVRegType lhsType, VMIVRegType rhsType, VMIVRegType resultType, + ArrayRef maskTypes, VMILayoutAttr requestedLayout, + bool preferredOnly, std::string *reason) { + auto fail = [&](const Twine &message) + -> FailureOr> { + if (reason) + *reason = message.str(); + return failure(); + }; + + if (lhsType.getElementCount() != rhsType.getElementCount() || + lhsType.getElementCount() != resultType.getElementCount()) + return fail("add-carry layout requires all data ports to share logical " + "lane count"); + if (lhsType.getElementType() != rhsType.getElementType() || + lhsType.getElementType() != resultType.getElementType()) + return fail("add-carry layout requires all data ports to share element " + "type"); + for (VMIMaskType maskType : maskTypes) { + if (maskType.getElementCount() != lhsType.getElementCount()) + return fail("add-carry layout requires all mask ports to share the data " + "lane count"); + if (maskType.getGranularity() != "b32") + return fail("add-carry layout requires b32 mask granularity"); + } + + int64_t numGroups = + requestedLayout && requestedLayout.isGroupSlots() + ? requestedLayout.getNumGroups() + : 0; + SmallVector facts; + for (const AddCarryLayoutPattern &pattern : kAddCarryLayoutPatterns) { + if (!matchesElementBitsPattern(pattern.dataBits, + lhsType.getElementType())) + continue; + if (llvm::any_of(maskTypes, [&](VMIMaskType maskType) { + return !matchesMaskGranularityPattern(pattern.maskGranularity, + maskType.getGranularity()); + })) + continue; + + VMILayoutAttr candidateLayout = materializeLayoutPattern( + lhsType.getContext(), pattern.layout, numGroups); + if (!candidateLayout || + (requestedLayout && candidateLayout != requestedLayout)) + continue; + + auto assignedDataType = VMIVRegType::get( + lhsType.getContext(), lhsType.getElementCount(), + lhsType.getElementType(), candidateLayout); + FailureOr dataArity = getVMIPhysicalArity(assignedDataType); + if (failed(dataArity) || *dataArity < 1) + continue; + + bool maskArityMatches = llvm::all_of(maskTypes, [&](VMIMaskType maskType) { + auto assignedMaskType = VMIMaskType::get( + maskType.getContext(), maskType.getElementCount(), + maskType.getGranularity(), candidateLayout); + FailureOr maskArity = getVMIPhysicalArity(assignedMaskType); + return succeeded(maskArity) && *maskArity == *dataArity; + }); + if (!maskArityMatches) + continue; + + facts.push_back(materializeAddCarryLayoutFact( + lhsType.getContext(), pattern, numGroups, *dataArity)); + if (preferredOnly) + break; + } + + if (facts.empty()) + return fail("add-carry ports do not match a legal layout table row"); + return facts; +} + +static FailureOr getAddCarryLayoutFactImpl( + VMIVRegType lhsType, VMIVRegType rhsType, VMIVRegType resultType, + ArrayRef maskTypes, std::string *reason) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr layout = lhsType.getLayoutAttr(); + if (!layout || rhsType.getLayoutAttr() != layout || + resultType.getLayoutAttr() != layout || + llvm::any_of(maskTypes, [&](VMIMaskType maskType) { + return maskType.getLayoutAttr() != layout; + })) + return fail("add-carry requires one assigned layout on every data, mask, " + "and carry port"); + + FailureOr> facts = + getAddCarryLayoutFactsForLayoutImpl(lhsType, rhsType, resultType, + maskTypes, layout, + /*preferredOnly=*/false, reason); + if (failed(facts)) + return failure(); + if (facts->size() != 1) + return fail("add-carry layout query produced ambiguous layout facts"); + return facts->front(); +} + static VMIVselrLayoutFact materializeVselrLayoutFact(MLIRContext *ctx, const VselrLayoutPattern &pattern) { @@ -2146,6 +2272,105 @@ VMILayoutSupport::getVdintlvLayoutFactForLayouts( reason); } +FailureOr +VMILayoutSupport::getPreferredVaddcLayoutFact(VMIVaddcOp op, + std::string *reason) const { + SmallVector maskTypes{ + cast(op.getMask().getType()), + cast(op.getCarry().getType())}; + FailureOr> facts = + getAddCarryLayoutFactsForLayoutImpl( + cast(op.getLhs().getType()), + cast(op.getRhs().getType()), + cast(op.getResult().getType()), maskTypes, + /*requestedLayout=*/{}, /*preferredOnly=*/true, reason); + if (failed(facts)) + return failure(); + return facts->front(); +} + +FailureOr +VMILayoutSupport::getPreferredVaddcsLayoutFact(VMIVaddcsOp op, + std::string *reason) const { + SmallVector maskTypes{ + cast(op.getCarryIn().getType()), + cast(op.getMask().getType()), + cast(op.getCarry().getType())}; + FailureOr> facts = + getAddCarryLayoutFactsForLayoutImpl( + cast(op.getLhs().getType()), + cast(op.getRhs().getType()), + cast(op.getResult().getType()), maskTypes, + /*requestedLayout=*/{}, /*preferredOnly=*/true, reason); + if (failed(facts)) + return failure(); + return facts->front(); +} + +FailureOr> +VMILayoutSupport::getVaddcLayoutFactsForLayout(VMIVaddcOp op, + VMILayoutAttr layout, + std::string *reason) const { + SmallVector maskTypes{ + cast(op.getMask().getType()), + cast(op.getCarry().getType())}; + return getAddCarryLayoutFactsForLayoutImpl( + cast(op.getLhs().getType()), + cast(op.getRhs().getType()), + cast(op.getResult().getType()), maskTypes, layout, + /*preferredOnly=*/false, reason); +} + +FailureOr> +VMILayoutSupport::getVaddcsLayoutFactsForLayout(VMIVaddcsOp op, + VMILayoutAttr layout, + std::string *reason) const { + SmallVector maskTypes{ + cast(op.getCarryIn().getType()), + cast(op.getMask().getType()), + cast(op.getCarry().getType())}; + return getAddCarryLayoutFactsForLayoutImpl( + cast(op.getLhs().getType()), + cast(op.getRhs().getType()), + cast(op.getResult().getType()), maskTypes, layout, + /*preferredOnly=*/false, reason); +} + +FailureOr +VMILayoutSupport::getVaddcLayoutFact(VMIVaddcOp op, + std::string *reason) const { + SmallVector maskTypes{ + cast(op.getMask().getType()), + cast(op.getCarry().getType())}; + return getAddCarryLayoutFactImpl( + cast(op.getLhs().getType()), + cast(op.getRhs().getType()), + cast(op.getResult().getType()), maskTypes, reason); +} + +FailureOr +VMILayoutSupport::getVaddcsLayoutFact(VMIVaddcsOp op, + std::string *reason) const { + SmallVector maskTypes{ + cast(op.getCarryIn().getType()), + cast(op.getMask().getType()), + cast(op.getCarry().getType())}; + return getAddCarryLayoutFactImpl( + cast(op.getLhs().getType()), + cast(op.getRhs().getType()), + cast(op.getResult().getType()), maskTypes, reason); +} + +LogicalResult VMILayoutSupport::getVaddcSupport(VMIVaddcOp op, + std::string *reason) const { + return getVaddcLayoutFact(op, reason); +} + +LogicalResult VMILayoutSupport::getVaddcsSupport(VMIVaddcsOp op, + std::string *reason) const { + return getVaddcsLayoutFact(op, reason); +} + FailureOr VMILayoutSupport::getLoadLayoutFact(VMIVRegType resultType, std::string *reason) const { diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 043329ce59..ce733e0e78 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -13008,70 +13008,14 @@ LogicalResult checkSupportedVmullShape(VMIVmullOp op, return success(); } -static LogicalResult -checkSupportedVMIAddCarryPorts(VMIVRegType lhsType, VMIVRegType rhsType, - VMIVRegType resultType, - ArrayRef maskTypes, - std::string *reason = nullptr) { - auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) - *reason = message.str(); - return failure(); - }; - - auto integerType = dyn_cast(lhsType.getElementType()); - if (!integerType || integerType.getWidth() != 32) - return fail("requires 32-bit integer data elements"); - if (lhsType != rhsType || lhsType != resultType) - return fail("requires matching lhs, rhs, and result VMI types"); - if (!lhsType.getLayoutAttr()) - return fail("requires assigned data layout"); - if (failed(checkSupportedMaskableVReg(lhsType))) - return fail("requires computable physical data parts"); - - FailureOr dataArity = getVMIPhysicalArity(lhsType); - if (failed(dataArity) || *dataArity < 1) - return fail("requires non-empty physical data parts"); - for (VMIMaskType maskType : maskTypes) { - if (maskType.getLayoutAttr() != lhsType.getLayoutAttr()) - return fail("requires all data and mask ports to share one layout"); - if (maskType.getGranularity() != "b32") - return fail("requires b32 mask granularity"); - FailureOr maskArity = getVMIPhysicalArity(maskType); - if (failed(maskArity) || *maskArity != *dataArity) - return fail("requires matching physical arity on data and mask ports"); - FailureOr physicalGranularity = - getVMIMaskPhysicalGranularity(maskType); - if (failed(physicalGranularity) || *physicalGranularity != "b32") - return fail("requires physical b32 mask parts"); - } - FailureOr lanesPerPart = getDataLanesPerPart(lhsType.getElementType()); - if (failed(lanesPerPart) || *lanesPerPart != 64) - return fail("requires 64-lane 32-bit data parts"); - return success(); -} - LogicalResult checkSupportedVMIAddcShape(VMIVaddcOp op, std::string *reason = nullptr) { - return checkSupportedVMIAddCarryPorts( - cast(op.getLhs().getType()), - cast(op.getRhs().getType()), - cast(op.getResult().getType()), - {cast(op.getMask().getType()), - cast(op.getCarry().getType())}, - reason); + return VMILayoutSupport().getVaddcSupport(op, reason); } LogicalResult checkSupportedVMIAddcsShape(VMIVaddcsOp op, std::string *reason = nullptr) { - return checkSupportedVMIAddCarryPorts( - cast(op.getLhs().getType()), - cast(op.getRhs().getType()), - cast(op.getResult().getType()), - {cast(op.getCarryIn().getType()), - cast(op.getMask().getType()), - cast(op.getCarry().getType())}, - reason); + return VMILayoutSupport().getVaddcsSupport(op, reason); } LogicalResult diff --git a/test/lit/vmi_new/vmi_carry_layout_support_invalid.pto b/test/lit/vmi_new/vmi_carry_layout_support_invalid.pto new file mode 100644 index 0000000000..994541ac06 --- /dev/null +++ b/test/lit/vmi_new/vmi_carry_layout_support_invalid.pto @@ -0,0 +1,55 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not pto-test-opt %s -split-input-file -vmi-to-vpto 2>&1 | FileCheck %s + +module { + func.func @vaddc_unregistered_group_slots( + %lhs: !pto.vmi.vreg<8xui32, #pto.vmi.layout>, + %rhs: !pto.vmi.vreg<8xui32, #pto.vmi.layout>, + %mask: !pto.vmi.mask<8xb32, #pto.vmi.layout>) + -> (!pto.vmi.vreg<8xui32, #pto.vmi.layout>, + !pto.vmi.mask<8xb32, #pto.vmi.layout>) { + %sum, %carry = pto.vmi.vaddc %lhs, %rhs, %mask + : !pto.vmi.vreg<8xui32, #pto.vmi.layout>, + !pto.vmi.vreg<8xui32, #pto.vmi.layout>, + !pto.vmi.mask<8xb32, #pto.vmi.layout> + -> !pto.vmi.vreg<8xui32, #pto.vmi.layout>, + !pto.vmi.mask<8xb32, #pto.vmi.layout> + return %sum, %carry + : !pto.vmi.vreg<8xui32, #pto.vmi.layout>, + !pto.vmi.mask<8xb32, #pto.vmi.layout> + } +} + +// CHECK: pto.vmi.vaddc requires matching 32-bit data and b32 mask parts +// CHECK-SAME: add-carry ports do not match a legal layout table row + +// ----- + +module { + func.func @vaddc_non_b32_physical_mask( + %lhs: !pto.vmi.vreg<64xui32, #pto.vmi.layout>, + %rhs: !pto.vmi.vreg<64xui32, #pto.vmi.layout>, + %mask: !pto.vmi.mask<64xb32, #pto.vmi.layout>) + -> (!pto.vmi.vreg<64xui32, #pto.vmi.layout>, + !pto.vmi.mask<64xb32, #pto.vmi.layout>) { + %sum, %carry = pto.vmi.vaddc %lhs, %rhs, %mask + : !pto.vmi.vreg<64xui32, #pto.vmi.layout>, + !pto.vmi.vreg<64xui32, #pto.vmi.layout>, + !pto.vmi.mask<64xb32, #pto.vmi.layout> + -> !pto.vmi.vreg<64xui32, #pto.vmi.layout>, + !pto.vmi.mask<64xb32, #pto.vmi.layout> + return %sum, %carry + : !pto.vmi.vreg<64xui32, #pto.vmi.layout>, + !pto.vmi.mask<64xb32, #pto.vmi.layout> + } +} + +// CHECK: pto.vmi.vaddc requires matching 32-bit data and b32 mask parts +// CHECK-SAME: add-carry ports do not match a legal layout table row diff --git a/test/lit/vmi_new/vmi_to_vpto_carry.pto b/test/lit/vmi_new/vmi_to_vpto_carry.pto index 803f5b105e..7f3b52ebe8 100644 --- a/test/lit/vmi_new/vmi_to_vpto_carry.pto +++ b/test/lit/vmi_new/vmi_to_vpto_carry.pto @@ -35,6 +35,23 @@ module { return %sum, %carry : !pto.vmi.vreg<128xui32>, !pto.vmi.mask<128xpred> } + + func.func @vmi_to_vpto_vaddc_deinterleaved( + %lhs: !pto.vmi.vreg<128xui32, #pto.vmi.layout>, + %rhs: !pto.vmi.vreg<128xui32, #pto.vmi.layout>, + %mask: !pto.vmi.mask<128xb32, #pto.vmi.layout>) + -> (!pto.vmi.vreg<128xui32, #pto.vmi.layout>, + !pto.vmi.mask<128xb32, #pto.vmi.layout>) { + %sum, %carry = pto.vmi.vaddc %lhs, %rhs, %mask + : !pto.vmi.vreg<128xui32, #pto.vmi.layout>, + !pto.vmi.vreg<128xui32, #pto.vmi.layout>, + !pto.vmi.mask<128xb32, #pto.vmi.layout> + -> !pto.vmi.vreg<128xui32, #pto.vmi.layout>, + !pto.vmi.mask<128xb32, #pto.vmi.layout> + return %sum, %carry + : !pto.vmi.vreg<128xui32, #pto.vmi.layout>, + !pto.vmi.mask<128xb32, #pto.vmi.layout> + } } // CHECK-LABEL: func.func @vmi_to_vpto_vaddc( @@ -45,5 +62,9 @@ module { // CHECK: %[[NEXT0:.*]], %[[NEXTCARRY0:.*]] = pto.vaddcs // CHECK: %[[NEXT1:.*]], %[[NEXTCARRY1:.*]] = pto.vaddcs // CHECK: return %[[NEXT0]], %[[NEXT1]], %[[NEXTCARRY0]], %[[NEXTCARRY1]] +// CHECK-LABEL: func.func @vmi_to_vpto_vaddc_deinterleaved( +// CHECK: %[[DEINT0:.*]], %[[DEINTCARRY0:.*]] = pto.vaddc +// CHECK: %[[DEINT1:.*]], %[[DEINTCARRY1:.*]] = pto.vaddc +// CHECK: return %[[DEINT0]], %[[DEINT1]], %[[DEINTCARRY0]], %[[DEINTCARRY1]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. From e0647ffa362cc45c944a58a31938ad4a0e25bdf0 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Mon, 10 Aug 2026 20:43:41 +0800 Subject: [PATCH 092/122] Revert "fix(vmi): centralize carry layout support" This reverts commit e5246354a198d2a28ee9bec037a2ae54866216e6. --- include/PTO/Transforms/VMILayoutSupport.h | 35 --- lib/PTO/Transforms/VMILayoutAssignment.cpp | 19 +- lib/PTO/Transforms/VMILayoutPropagation.cpp | 60 +---- lib/PTO/Transforms/VMILayoutSupport.cpp | 225 ------------------ lib/PTO/Transforms/VMIToVPTO.cpp | 60 ++++- .../vmi_carry_layout_support_invalid.pto | 55 ----- test/lit/vmi_new/vmi_to_vpto_carry.pto | 21 -- 7 files changed, 61 insertions(+), 414 deletions(-) delete mode 100644 test/lit/vmi_new/vmi_carry_layout_support_invalid.pto diff --git a/include/PTO/Transforms/VMILayoutSupport.h b/include/PTO/Transforms/VMILayoutSupport.h index bf9fb5d085..179640e9c9 100644 --- a/include/PTO/Transforms/VMILayoutSupport.h +++ b/include/PTO/Transforms/VMILayoutSupport.h @@ -104,11 +104,6 @@ struct VMIInterleaveLayoutFact { int64_t lanesPerPart = 0; }; -struct VMIAddCarryLayoutFact { - VMILayoutAttr layout; - int64_t physicalArity = 0; -}; - struct VMIBitcastLayoutFact { VMILayoutAttr sourceLayout; VMILayoutAttr resultLayout; @@ -321,30 +316,6 @@ class VMILayoutSupport { VMIVRegType lowType, VMIVRegType highType, std::string *reason = nullptr) const; - FailureOr - getPreferredVaddcLayoutFact(VMIVaddcOp op, - std::string *reason = nullptr) const; - - FailureOr - getPreferredVaddcsLayoutFact(VMIVaddcsOp op, - std::string *reason = nullptr) const; - - FailureOr> - getVaddcLayoutFactsForLayout(VMIVaddcOp op, VMILayoutAttr layout, - std::string *reason = nullptr) const; - - FailureOr> - getVaddcsLayoutFactsForLayout(VMIVaddcsOp op, VMILayoutAttr layout, - std::string *reason = nullptr) const; - - FailureOr - getVaddcLayoutFact(VMIVaddcOp op, - std::string *reason = nullptr) const; - - FailureOr - getVaddcsLayoutFact(VMIVaddcsOp op, - std::string *reason = nullptr) const; - FailureOr getGroupSlotLoadLayoutFact(VMIVRegType resultType, int64_t numGroups, std::string *reason = nullptr) const; @@ -439,12 +410,6 @@ class VMILayoutSupport { LogicalResult getVselrSupport(VMIVselrOp op, std::string *reason = nullptr) const; - LogicalResult getVaddcSupport(VMIVaddcOp op, - std::string *reason = nullptr) const; - - LogicalResult getVaddcsSupport(VMIVaddcsOp op, - std::string *reason = nullptr) const; - LogicalResult getGroupReduceAddFSupport(VMIGroupReduceAddFOp op, std::string *reason = nullptr) const; diff --git a/lib/PTO/Transforms/VMILayoutAssignment.cpp b/lib/PTO/Transforms/VMILayoutAssignment.cpp index 7710c154ce..49555602fc 100644 --- a/lib/PTO/Transforms/VMILayoutAssignment.cpp +++ b/lib/PTO/Transforms/VMILayoutAssignment.cpp @@ -656,31 +656,16 @@ struct LayoutSolver { return WalkResult::advance(); } if (auto addc = dyn_cast(op)) { - VMILayoutSupport supports; - std::string reason; - if (failed(supports.getPreferredVaddcLayoutFact(addc, &reason))) { - addc.emitError() << kVMIDiagLayoutContractPrefix << reason; - return WalkResult::interrupt(); - } if (failed(constrainElementwiseBinary(addc.getLhsMutable(), addc.getRhsMutable(), - addc.getResult(), op)) || - failed(uniteMask(addc.getMask(), addc.getCarry(), op))) + addc.getResult(), op))) return WalkResult::interrupt(); return WalkResult::advance(); } if (auto addcs = dyn_cast(op)) { - VMILayoutSupport supports; - std::string reason; - if (failed(supports.getPreferredVaddcsLayoutFact(addcs, &reason))) { - addcs.emitError() << kVMIDiagLayoutContractPrefix << reason; - return WalkResult::interrupt(); - } if (failed(constrainElementwiseBinary(addcs.getLhsMutable(), addcs.getRhsMutable(), - addcs.getResult(), op)) || - failed(uniteMask(addcs.getCarryIn(), addcs.getMask(), op)) || - failed(uniteMask(addcs.getMask(), addcs.getCarry(), op))) + addcs.getResult(), op))) return WalkResult::interrupt(); return WalkResult::advance(); } diff --git a/lib/PTO/Transforms/VMILayoutPropagation.cpp b/lib/PTO/Transforms/VMILayoutPropagation.cpp index 78d42b7e5f..d3004656f9 100644 --- a/lib/PTO/Transforms/VMILayoutPropagation.cpp +++ b/lib/PTO/Transforms/VMILayoutPropagation.cpp @@ -170,6 +170,7 @@ class VMILayoutMaterializationTransfer final { static bool isSameLayoutOp(Operation *op) { return isa> - query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, - const VMILayoutPropagator &propagator, - OpOperand *changedOperand) const override { - if (auto addc = dyn_cast(op)) - return queryVaddc(addc, changedLayout); - if (auto addcs = dyn_cast(op)) - return queryVaddcs(addcs, changedLayout); - return failure(); - } - -private: - FailureOr> - queryVaddc(VMIVaddcOp op, VMILayoutAttr changedLayout) const { - VMILayoutSupport supports; - FailureOr> facts = - supports.getVaddcLayoutFactsForLayout(op, changedLayout); - if (failed(facts) || facts->empty()) - return failure(); - - SmallVector relations; - for (const VMIAddCarryLayoutFact &fact : *facts) { - relations.push_back(makeRelation(SmallVector{ - operandFact(op.getLhsMutable(), fact.layout), - operandFact(op.getRhsMutable(), fact.layout), - operandFact(op.getMaskMutable(), fact.layout), - valueFact(op.getResult(), fact.layout), - valueFact(op.getCarry(), fact.layout)})); - } - return relations; - } - - FailureOr> - queryVaddcs(VMIVaddcsOp op, VMILayoutAttr changedLayout) const { - VMILayoutSupport supports; - FailureOr> facts = - supports.getVaddcsLayoutFactsForLayout(op, changedLayout); - if (failed(facts) || facts->empty()) - return failure(); - - SmallVector relations; - for (const VMIAddCarryLayoutFact &fact : *facts) { - relations.push_back(makeRelation(SmallVector{ - operandFact(op.getLhsMutable(), fact.layout), - operandFact(op.getRhsMutable(), fact.layout), - operandFact(op.getCarryInMutable(), fact.layout), - operandFact(op.getMaskMutable(), fact.layout), - valueFact(op.getResult(), fact.layout), - valueFact(op.getCarry(), fact.layout)})); - } - return relations; - } -}; - class VMIInterleaveTransfer final : public VMILayoutTransfer { public: FailureOr> @@ -889,7 +834,6 @@ const VMILayoutTransfer *getTransfer(Operation *op) { static VMIGroupSlotLoadTransfer groupSlotLoadTransfer; static VMIGroupBroadcastLoadTransfer groupBroadcastLoadTransfer; static VMIGroupBroadcastTransfer groupBroadcastTransfer; - static VMIAddCarryTransfer addCarryTransfer; static VMIInterleaveTransfer interleaveTransfer; static VMIGatherTransfer gatherTransfer; static VMIStoreTransfer storeTransfer; @@ -918,8 +862,6 @@ const VMILayoutTransfer *getTransfer(Operation *op) { return &groupBroadcastLoadTransfer; if (isa(op)) return &groupBroadcastTransfer; - if (isa(op)) - return &addCarryTransfer; if (isa(op)) return &interleaveTransfer; if (isa(op)) diff --git a/lib/PTO/Transforms/VMILayoutSupport.cpp b/lib/PTO/Transforms/VMILayoutSupport.cpp index 0b6372f56f..4ec44e2e7f 100644 --- a/lib/PTO/Transforms/VMILayoutSupport.cpp +++ b/lib/PTO/Transforms/VMILayoutSupport.cpp @@ -451,12 +451,6 @@ struct InterleaveLayoutPattern { LayoutPattern highLayout; }; -struct AddCarryLayoutPattern { - ElementBitsPattern dataBits; - MaskGranularityPattern maskGranularity; - LayoutPattern layout; -}; - static constexpr PreferredCastLayoutPattern kPreferredCastLayoutPatterns[] = { // Exact rows override the default legal relation for small shapes where the // compact lane-stride form is the natural cast layout. @@ -610,13 +604,6 @@ static constexpr InterleaveLayoutPattern kVintlvLayoutPatterns[] = { {bits<8>(), chunk<1>(), 1, ls(4), ls(4), ls(4), ls(4), ls(4)}, }; -static constexpr AddCarryLayoutPattern kAddCarryLayoutPatterns[] = { - {bits<32>(), mb32(), c()}, {bits<32>(), mb32(), d(2)}, - {bits<32>(), mb32(), d(4)}, {bits<32>(), mb32(), bd(2)}, - {bits<32>(), mb32(), bd(4)}, {bits<32>(), mb32(), gs(1)}, - {bits<32>(), mb32(), gs(8)}, -}; - struct DenseMemoryLayoutPattern { ElementBitsPattern elementBits; LayoutPattern layout; @@ -1135,119 +1122,6 @@ static VMIInterleaveLayoutFact materializeInterleaveLayoutFact( return fact; } -static VMIAddCarryLayoutFact materializeAddCarryLayoutFact( - MLIRContext *ctx, const AddCarryLayoutPattern &pattern, int64_t numGroups, - int64_t physicalArity) { - return VMIAddCarryLayoutFact{ - materializeLayoutPattern(ctx, pattern.layout, numGroups), physicalArity}; -} - -static FailureOr> -getAddCarryLayoutFactsForLayoutImpl( - VMIVRegType lhsType, VMIVRegType rhsType, VMIVRegType resultType, - ArrayRef maskTypes, VMILayoutAttr requestedLayout, - bool preferredOnly, std::string *reason) { - auto fail = [&](const Twine &message) - -> FailureOr> { - if (reason) - *reason = message.str(); - return failure(); - }; - - if (lhsType.getElementCount() != rhsType.getElementCount() || - lhsType.getElementCount() != resultType.getElementCount()) - return fail("add-carry layout requires all data ports to share logical " - "lane count"); - if (lhsType.getElementType() != rhsType.getElementType() || - lhsType.getElementType() != resultType.getElementType()) - return fail("add-carry layout requires all data ports to share element " - "type"); - for (VMIMaskType maskType : maskTypes) { - if (maskType.getElementCount() != lhsType.getElementCount()) - return fail("add-carry layout requires all mask ports to share the data " - "lane count"); - if (maskType.getGranularity() != "b32") - return fail("add-carry layout requires b32 mask granularity"); - } - - int64_t numGroups = - requestedLayout && requestedLayout.isGroupSlots() - ? requestedLayout.getNumGroups() - : 0; - SmallVector facts; - for (const AddCarryLayoutPattern &pattern : kAddCarryLayoutPatterns) { - if (!matchesElementBitsPattern(pattern.dataBits, - lhsType.getElementType())) - continue; - if (llvm::any_of(maskTypes, [&](VMIMaskType maskType) { - return !matchesMaskGranularityPattern(pattern.maskGranularity, - maskType.getGranularity()); - })) - continue; - - VMILayoutAttr candidateLayout = materializeLayoutPattern( - lhsType.getContext(), pattern.layout, numGroups); - if (!candidateLayout || - (requestedLayout && candidateLayout != requestedLayout)) - continue; - - auto assignedDataType = VMIVRegType::get( - lhsType.getContext(), lhsType.getElementCount(), - lhsType.getElementType(), candidateLayout); - FailureOr dataArity = getVMIPhysicalArity(assignedDataType); - if (failed(dataArity) || *dataArity < 1) - continue; - - bool maskArityMatches = llvm::all_of(maskTypes, [&](VMIMaskType maskType) { - auto assignedMaskType = VMIMaskType::get( - maskType.getContext(), maskType.getElementCount(), - maskType.getGranularity(), candidateLayout); - FailureOr maskArity = getVMIPhysicalArity(assignedMaskType); - return succeeded(maskArity) && *maskArity == *dataArity; - }); - if (!maskArityMatches) - continue; - - facts.push_back(materializeAddCarryLayoutFact( - lhsType.getContext(), pattern, numGroups, *dataArity)); - if (preferredOnly) - break; - } - - if (facts.empty()) - return fail("add-carry ports do not match a legal layout table row"); - return facts; -} - -static FailureOr getAddCarryLayoutFactImpl( - VMIVRegType lhsType, VMIVRegType rhsType, VMIVRegType resultType, - ArrayRef maskTypes, std::string *reason) { - auto fail = [&](const Twine &message) -> FailureOr { - if (reason) - *reason = message.str(); - return failure(); - }; - - VMILayoutAttr layout = lhsType.getLayoutAttr(); - if (!layout || rhsType.getLayoutAttr() != layout || - resultType.getLayoutAttr() != layout || - llvm::any_of(maskTypes, [&](VMIMaskType maskType) { - return maskType.getLayoutAttr() != layout; - })) - return fail("add-carry requires one assigned layout on every data, mask, " - "and carry port"); - - FailureOr> facts = - getAddCarryLayoutFactsForLayoutImpl(lhsType, rhsType, resultType, - maskTypes, layout, - /*preferredOnly=*/false, reason); - if (failed(facts)) - return failure(); - if (facts->size() != 1) - return fail("add-carry layout query produced ambiguous layout facts"); - return facts->front(); -} - static VMIVselrLayoutFact materializeVselrLayoutFact(MLIRContext *ctx, const VselrLayoutPattern &pattern) { @@ -2272,105 +2146,6 @@ VMILayoutSupport::getVdintlvLayoutFactForLayouts( reason); } -FailureOr -VMILayoutSupport::getPreferredVaddcLayoutFact(VMIVaddcOp op, - std::string *reason) const { - SmallVector maskTypes{ - cast(op.getMask().getType()), - cast(op.getCarry().getType())}; - FailureOr> facts = - getAddCarryLayoutFactsForLayoutImpl( - cast(op.getLhs().getType()), - cast(op.getRhs().getType()), - cast(op.getResult().getType()), maskTypes, - /*requestedLayout=*/{}, /*preferredOnly=*/true, reason); - if (failed(facts)) - return failure(); - return facts->front(); -} - -FailureOr -VMILayoutSupport::getPreferredVaddcsLayoutFact(VMIVaddcsOp op, - std::string *reason) const { - SmallVector maskTypes{ - cast(op.getCarryIn().getType()), - cast(op.getMask().getType()), - cast(op.getCarry().getType())}; - FailureOr> facts = - getAddCarryLayoutFactsForLayoutImpl( - cast(op.getLhs().getType()), - cast(op.getRhs().getType()), - cast(op.getResult().getType()), maskTypes, - /*requestedLayout=*/{}, /*preferredOnly=*/true, reason); - if (failed(facts)) - return failure(); - return facts->front(); -} - -FailureOr> -VMILayoutSupport::getVaddcLayoutFactsForLayout(VMIVaddcOp op, - VMILayoutAttr layout, - std::string *reason) const { - SmallVector maskTypes{ - cast(op.getMask().getType()), - cast(op.getCarry().getType())}; - return getAddCarryLayoutFactsForLayoutImpl( - cast(op.getLhs().getType()), - cast(op.getRhs().getType()), - cast(op.getResult().getType()), maskTypes, layout, - /*preferredOnly=*/false, reason); -} - -FailureOr> -VMILayoutSupport::getVaddcsLayoutFactsForLayout(VMIVaddcsOp op, - VMILayoutAttr layout, - std::string *reason) const { - SmallVector maskTypes{ - cast(op.getCarryIn().getType()), - cast(op.getMask().getType()), - cast(op.getCarry().getType())}; - return getAddCarryLayoutFactsForLayoutImpl( - cast(op.getLhs().getType()), - cast(op.getRhs().getType()), - cast(op.getResult().getType()), maskTypes, layout, - /*preferredOnly=*/false, reason); -} - -FailureOr -VMILayoutSupport::getVaddcLayoutFact(VMIVaddcOp op, - std::string *reason) const { - SmallVector maskTypes{ - cast(op.getMask().getType()), - cast(op.getCarry().getType())}; - return getAddCarryLayoutFactImpl( - cast(op.getLhs().getType()), - cast(op.getRhs().getType()), - cast(op.getResult().getType()), maskTypes, reason); -} - -FailureOr -VMILayoutSupport::getVaddcsLayoutFact(VMIVaddcsOp op, - std::string *reason) const { - SmallVector maskTypes{ - cast(op.getCarryIn().getType()), - cast(op.getMask().getType()), - cast(op.getCarry().getType())}; - return getAddCarryLayoutFactImpl( - cast(op.getLhs().getType()), - cast(op.getRhs().getType()), - cast(op.getResult().getType()), maskTypes, reason); -} - -LogicalResult VMILayoutSupport::getVaddcSupport(VMIVaddcOp op, - std::string *reason) const { - return getVaddcLayoutFact(op, reason); -} - -LogicalResult VMILayoutSupport::getVaddcsSupport(VMIVaddcsOp op, - std::string *reason) const { - return getVaddcsLayoutFact(op, reason); -} - FailureOr VMILayoutSupport::getLoadLayoutFact(VMIVRegType resultType, std::string *reason) const { diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index ce733e0e78..043329ce59 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -13008,14 +13008,70 @@ LogicalResult checkSupportedVmullShape(VMIVmullOp op, return success(); } +static LogicalResult +checkSupportedVMIAddCarryPorts(VMIVRegType lhsType, VMIVRegType rhsType, + VMIVRegType resultType, + ArrayRef maskTypes, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto integerType = dyn_cast(lhsType.getElementType()); + if (!integerType || integerType.getWidth() != 32) + return fail("requires 32-bit integer data elements"); + if (lhsType != rhsType || lhsType != resultType) + return fail("requires matching lhs, rhs, and result VMI types"); + if (!lhsType.getLayoutAttr()) + return fail("requires assigned data layout"); + if (failed(checkSupportedMaskableVReg(lhsType))) + return fail("requires computable physical data parts"); + + FailureOr dataArity = getVMIPhysicalArity(lhsType); + if (failed(dataArity) || *dataArity < 1) + return fail("requires non-empty physical data parts"); + for (VMIMaskType maskType : maskTypes) { + if (maskType.getLayoutAttr() != lhsType.getLayoutAttr()) + return fail("requires all data and mask ports to share one layout"); + if (maskType.getGranularity() != "b32") + return fail("requires b32 mask granularity"); + FailureOr maskArity = getVMIPhysicalArity(maskType); + if (failed(maskArity) || *maskArity != *dataArity) + return fail("requires matching physical arity on data and mask ports"); + FailureOr physicalGranularity = + getVMIMaskPhysicalGranularity(maskType); + if (failed(physicalGranularity) || *physicalGranularity != "b32") + return fail("requires physical b32 mask parts"); + } + FailureOr lanesPerPart = getDataLanesPerPart(lhsType.getElementType()); + if (failed(lanesPerPart) || *lanesPerPart != 64) + return fail("requires 64-lane 32-bit data parts"); + return success(); +} + LogicalResult checkSupportedVMIAddcShape(VMIVaddcOp op, std::string *reason = nullptr) { - return VMILayoutSupport().getVaddcSupport(op, reason); + return checkSupportedVMIAddCarryPorts( + cast(op.getLhs().getType()), + cast(op.getRhs().getType()), + cast(op.getResult().getType()), + {cast(op.getMask().getType()), + cast(op.getCarry().getType())}, + reason); } LogicalResult checkSupportedVMIAddcsShape(VMIVaddcsOp op, std::string *reason = nullptr) { - return VMILayoutSupport().getVaddcsSupport(op, reason); + return checkSupportedVMIAddCarryPorts( + cast(op.getLhs().getType()), + cast(op.getRhs().getType()), + cast(op.getResult().getType()), + {cast(op.getCarryIn().getType()), + cast(op.getMask().getType()), + cast(op.getCarry().getType())}, + reason); } LogicalResult diff --git a/test/lit/vmi_new/vmi_carry_layout_support_invalid.pto b/test/lit/vmi_new/vmi_carry_layout_support_invalid.pto deleted file mode 100644 index 994541ac06..0000000000 --- a/test/lit/vmi_new/vmi_carry_layout_support_invalid.pto +++ /dev/null @@ -1,55 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -// RUN: not pto-test-opt %s -split-input-file -vmi-to-vpto 2>&1 | FileCheck %s - -module { - func.func @vaddc_unregistered_group_slots( - %lhs: !pto.vmi.vreg<8xui32, #pto.vmi.layout>, - %rhs: !pto.vmi.vreg<8xui32, #pto.vmi.layout>, - %mask: !pto.vmi.mask<8xb32, #pto.vmi.layout>) - -> (!pto.vmi.vreg<8xui32, #pto.vmi.layout>, - !pto.vmi.mask<8xb32, #pto.vmi.layout>) { - %sum, %carry = pto.vmi.vaddc %lhs, %rhs, %mask - : !pto.vmi.vreg<8xui32, #pto.vmi.layout>, - !pto.vmi.vreg<8xui32, #pto.vmi.layout>, - !pto.vmi.mask<8xb32, #pto.vmi.layout> - -> !pto.vmi.vreg<8xui32, #pto.vmi.layout>, - !pto.vmi.mask<8xb32, #pto.vmi.layout> - return %sum, %carry - : !pto.vmi.vreg<8xui32, #pto.vmi.layout>, - !pto.vmi.mask<8xb32, #pto.vmi.layout> - } -} - -// CHECK: pto.vmi.vaddc requires matching 32-bit data and b32 mask parts -// CHECK-SAME: add-carry ports do not match a legal layout table row - -// ----- - -module { - func.func @vaddc_non_b32_physical_mask( - %lhs: !pto.vmi.vreg<64xui32, #pto.vmi.layout>, - %rhs: !pto.vmi.vreg<64xui32, #pto.vmi.layout>, - %mask: !pto.vmi.mask<64xb32, #pto.vmi.layout>) - -> (!pto.vmi.vreg<64xui32, #pto.vmi.layout>, - !pto.vmi.mask<64xb32, #pto.vmi.layout>) { - %sum, %carry = pto.vmi.vaddc %lhs, %rhs, %mask - : !pto.vmi.vreg<64xui32, #pto.vmi.layout>, - !pto.vmi.vreg<64xui32, #pto.vmi.layout>, - !pto.vmi.mask<64xb32, #pto.vmi.layout> - -> !pto.vmi.vreg<64xui32, #pto.vmi.layout>, - !pto.vmi.mask<64xb32, #pto.vmi.layout> - return %sum, %carry - : !pto.vmi.vreg<64xui32, #pto.vmi.layout>, - !pto.vmi.mask<64xb32, #pto.vmi.layout> - } -} - -// CHECK: pto.vmi.vaddc requires matching 32-bit data and b32 mask parts -// CHECK-SAME: add-carry ports do not match a legal layout table row diff --git a/test/lit/vmi_new/vmi_to_vpto_carry.pto b/test/lit/vmi_new/vmi_to_vpto_carry.pto index 7f3b52ebe8..803f5b105e 100644 --- a/test/lit/vmi_new/vmi_to_vpto_carry.pto +++ b/test/lit/vmi_new/vmi_to_vpto_carry.pto @@ -35,23 +35,6 @@ module { return %sum, %carry : !pto.vmi.vreg<128xui32>, !pto.vmi.mask<128xpred> } - - func.func @vmi_to_vpto_vaddc_deinterleaved( - %lhs: !pto.vmi.vreg<128xui32, #pto.vmi.layout>, - %rhs: !pto.vmi.vreg<128xui32, #pto.vmi.layout>, - %mask: !pto.vmi.mask<128xb32, #pto.vmi.layout>) - -> (!pto.vmi.vreg<128xui32, #pto.vmi.layout>, - !pto.vmi.mask<128xb32, #pto.vmi.layout>) { - %sum, %carry = pto.vmi.vaddc %lhs, %rhs, %mask - : !pto.vmi.vreg<128xui32, #pto.vmi.layout>, - !pto.vmi.vreg<128xui32, #pto.vmi.layout>, - !pto.vmi.mask<128xb32, #pto.vmi.layout> - -> !pto.vmi.vreg<128xui32, #pto.vmi.layout>, - !pto.vmi.mask<128xb32, #pto.vmi.layout> - return %sum, %carry - : !pto.vmi.vreg<128xui32, #pto.vmi.layout>, - !pto.vmi.mask<128xb32, #pto.vmi.layout> - } } // CHECK-LABEL: func.func @vmi_to_vpto_vaddc( @@ -62,9 +45,5 @@ module { // CHECK: %[[NEXT0:.*]], %[[NEXTCARRY0:.*]] = pto.vaddcs // CHECK: %[[NEXT1:.*]], %[[NEXTCARRY1:.*]] = pto.vaddcs // CHECK: return %[[NEXT0]], %[[NEXT1]], %[[NEXTCARRY0]], %[[NEXTCARRY1]] -// CHECK-LABEL: func.func @vmi_to_vpto_vaddc_deinterleaved( -// CHECK: %[[DEINT0:.*]], %[[DEINTCARRY0:.*]] = pto.vaddc -// CHECK: %[[DEINT1:.*]], %[[DEINTCARRY1:.*]] = pto.vaddc -// CHECK: return %[[DEINT0]], %[[DEINT1]], %[[DEINTCARRY0]], %[[DEINTCARRY1]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. From a847236ee325b83395553bb914cca2014e5087b1 Mon Sep 17 00:00:00 2001 From: qukelin Date: Tue, 11 Aug 2026 02:32:00 +0800 Subject: [PATCH 093/122] feat(vpto): support scalar f16 fmin and fmax --- docs/isa/micro-isa/17-simt.md | 14 ++++++++------ include/PTO/IR/VPTOOps.td | 17 +++++++++-------- lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp | 8 ++++---- lib/PTO/Transforms/VPTOLLVMEmitter.cpp | 8 ++++---- ptodsl/docs/user_guide/13-simt-micro-ops.md | 6 +++++- ptodsl/tests/test_jit_compile.py | 11 +++++++++++ .../simt_lowlevel_float_math_vpto_llvm.pto | 7 +++++++ .../simt/simt-float-convert-core/compare.py | 10 ++++++++-- .../simt/simt-float-convert-core/golden.py | 4 ++++ .../simt/simt-float-convert-core/kernel.pto | 19 +++++++++++++++++++ 10 files changed, 79 insertions(+), 25 deletions(-) diff --git a/docs/isa/micro-isa/17-simt.md b/docs/isa/micro-isa/17-simt.md index 2068f61cb7..73db6575a1 100644 --- a/docs/isa/micro-isa/17-simt.md +++ b/docs/isa/micro-isa/17-simt.md @@ -880,9 +880,10 @@ else: - **semantics:** Return the floating minimum of `%a` and `%b`. - **inputs:** `%a` and `%b` have the same type. - **outputs:** One value with the same type as the inputs. -- **constraints and limitations:** `T` is `f32`, `bf16`, `vector<2xf16>`, or - `vector<2xbf16>`. For vector types, the minimum is computed element-wise. NaN - handling follows the target floating-point minimum rule. +- **constraints and limitations:** `T` is `f16`, `f32`, `bf16`, + `vector<2xf16>`, or `vector<2xbf16>`. For vector types, the minimum is + computed element-wise. NaN handling follows the target floating-point + minimum rule. ### `pto.fmax` @@ -890,9 +891,10 @@ else: - **semantics:** Return the floating maximum of `%a` and `%b`. - **inputs:** `%a` and `%b` have the same type. - **outputs:** One value with the same type as the inputs. -- **constraints and limitations:** `T` is `f32`, `bf16`, `vector<2xf16>`, or - `vector<2xbf16>`. For vector types, the maximum is computed element-wise. NaN - handling follows the target floating-point maximum rule. +- **constraints and limitations:** `T` is `f16`, `f32`, `bf16`, + `vector<2xf16>`, or `vector<2xbf16>`. For vector types, the maximum is + computed element-wise. NaN handling follows the target floating-point + maximum rule. ### `pto.fma` diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index d8bfc8b72a..bd1ab9d08b 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -124,8 +124,9 @@ def PTO_SimtUnaryFloatValueType : AnyTypeOf<[F16, F32, BF16, PTO_V2F16Type, PTO_ "f16, f32, bf16, vector<2xf16> or vector<2xbf16>">; def PTO_SimtUnaryF16F32ValueType : AnyTypeOf<[F16, F32, PTO_V2F16Type], "f16, f32 or vector<2xf16>">; -def PTO_SimtBinaryF32BF16ValueType : AnyTypeOf<[F32, BF16, PTO_V2F16Type, PTO_V2BF16Type], - "f32, bf16, vector<2xf16> or vector<2xbf16>">; +def PTO_SimtBinaryFloatValueType : AnyTypeOf< + [F16, F32, BF16, PTO_V2F16Type, PTO_V2BF16Type], + "f16, f32, bf16, vector<2xf16> or vector<2xbf16>">; def PTO_SimtBinaryF16F32ValueType : AnyTypeOf<[F16, F32, PTO_V2F16Type], "f16, f32 or vector<2xf16>">; def PTO_SimtConvertValueType : AnyTypeOf<[I32, I64, F16, BF16, F32, @@ -924,13 +925,13 @@ def PTO_FloorOp : PTO_UnaryFloatScalarOp<"floor">; def PTO_RintOp : PTO_UnaryFloatScalarOp<"rint">; def PTO_RoundOp : PTO_UnaryFloatScalarOp<"round">; -class PTO_BinaryF32BF16ScalarOp +class PTO_BinaryFloatValueOp : PTO_SimtOp]> { let arguments = (ins - PTO_SimtBinaryF32BF16ValueType:$lhs, - PTO_SimtBinaryF32BF16ValueType:$rhs + PTO_SimtBinaryFloatValueType:$lhs, + PTO_SimtBinaryFloatValueType:$rhs ); - let results = (outs PTO_SimtBinaryF32BF16ValueType:$result); + let results = (outs PTO_SimtBinaryFloatValueType:$result); let assemblyFormat = [{ $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result) @@ -950,8 +951,8 @@ class PTO_BinaryF16F32ScalarOp }]; } -def PTO_FMinOp : PTO_BinaryF32BF16ScalarOp<"fmin">; -def PTO_FMaxOp : PTO_BinaryF32BF16ScalarOp<"fmax">; +def PTO_FMinOp : PTO_BinaryFloatValueOp<"fmin">; +def PTO_FMaxOp : PTO_BinaryFloatValueOp<"fmax">; def PTO_PowOp : PTO_BinaryF16F32ScalarOp<"pow">; def PTO_FmaOp diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index 767b409633..9ce5c0c3c5 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -3051,8 +3051,8 @@ template <> FailureOr buildBinaryScalarMathCallee(MLIRContext *context, Type valueType) { std::string elem = getLLVMFloatBuiltinFragment(valueType); - if (elem != "f32" && elem != "bf16" && elem != "v2f16" && - elem != "v2bf16") + if (elem != "f16" && elem != "f32" && elem != "bf16" && + elem != "v2f16" && elem != "v2bf16") return failure(); return StringAttr::get(context, "llvm.minnum." + elem).getValue(); } @@ -3061,8 +3061,8 @@ template <> FailureOr buildBinaryScalarMathCallee(MLIRContext *context, Type valueType) { std::string elem = getLLVMFloatBuiltinFragment(valueType); - if (elem != "f32" && elem != "bf16" && elem != "v2f16" && - elem != "v2bf16") + if (elem != "f16" && elem != "f32" && elem != "bf16" && + elem != "v2f16" && elem != "v2bf16") return failure(); return StringAttr::get(context, "llvm.maxnum." + elem).getValue(); } diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index e98e091dee..68356be08f 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -3079,8 +3079,8 @@ template <> FailureOr buildBinaryScalarMathCallee(MLIRContext *context, Type valueType) { std::string elem = getLLVMFloatBuiltinFragment(valueType); - if (elem != "f32" && elem != "bf16" && elem != "v2f16" && - elem != "v2bf16") + if (elem != "f16" && elem != "f32" && elem != "bf16" && + elem != "v2f16" && elem != "v2bf16") return failure(); return StringAttr::get(context, "llvm.minnum." + elem).getValue(); } @@ -3089,8 +3089,8 @@ template <> FailureOr buildBinaryScalarMathCallee(MLIRContext *context, Type valueType) { std::string elem = getLLVMFloatBuiltinFragment(valueType); - if (elem != "f32" && elem != "bf16" && elem != "v2f16" && - elem != "v2bf16") + if (elem != "f16" && elem != "f32" && elem != "bf16" && + elem != "v2f16" && elem != "v2bf16") return failure(); return StringAttr::get(context, "llvm.maxnum." + elem).getValue(); } diff --git a/ptodsl/docs/user_guide/13-simt-micro-ops.md b/ptodsl/docs/user_guide/13-simt-micro-ops.md index eec1f3bfbc..8cb244b948 100644 --- a/ptodsl/docs/user_guide/13-simt-micro-ops.md +++ b/ptodsl/docs/user_guide/13-simt-micro-ops.md @@ -461,7 +461,11 @@ def simt_ops_integer_math_probe(dst: pto.ptr(pto.i32, "gm")): **Description**: Performs SIMT floating-point math. These functions are VPTO SIMT micro-ops and are distinct from the generic scalar helpers in Chapter 6. -**Parameters**: PTO floating-point scalar operands. +`pto.fmin` and `pto.fmax` accept `f16`, `f32`, `bf16`, `vector<2xf16>`, and +`vector<2xbf16>` operands. + +**Parameters**: PTO floating-point scalar or packed operands supported by the +selected operation. **Returns**: PTO scalar with the same type as the input value. diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index c2d6bdc708..3f77d444b0 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -891,6 +891,7 @@ def simt_collective_math_probe(): pto.mul_i32toi64(lane, lane, signedness="unsigned") as_f32 = pto.convert(lane, pto.f32, rounding="r", saturation="nosat", signedness="signed") + as_f16 = pto.const(1.0, dtype=pto.f16) pto.convert(as_f32, pto.i32, rounding="z", saturation="sat", signedness="signed") pto.absf(as_f32) pto.sqrt(as_f32) @@ -903,6 +904,8 @@ def simt_collective_math_probe(): pto.round(as_f32) pto.fmin(as_f32, as_f32) pto.fmax(as_f32, as_f32) + pto.fmin(as_f16, as_f16) + pto.fmax(as_f16, as_f16) pto.fma(as_f32, as_f32, as_f32) @@ -5639,6 +5642,14 @@ def _enter_inline_simt_with_resource_attr(): "pto.resume", ): expect(op_name in simt_full_text, f"full SIMT surface should contain {op_name}") + expect( + re.search(r"pto\.fmin .* : f16, f16 -> f16", simt_full_text) is not None, + "full SIMT surface should accept scalar f16 pto.fmin", + ) + expect( + re.search(r"pto\.fmax .* : f16, f16 -> f16", simt_full_text) is not None, + "full SIMT surface should accept scalar f16 pto.fmax", + ) for fp8_vec in ( "vector<4xf8E4M3FN>", "vector<8xf8E4M3FN>", diff --git a/test/lit/vpto/simt_lowlevel_float_math_vpto_llvm.pto b/test/lit/vpto/simt_lowlevel_float_math_vpto_llvm.pto index 0238af35d8..4d634c5af7 100644 --- a/test/lit/vpto/simt_lowlevel_float_math_vpto_llvm.pto +++ b/test/lit/vpto/simt_lowlevel_float_math_vpto_llvm.pto @@ -7,6 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ( mkdir -p %T && ptoas --pto-arch=a5 --pto-backend=vpto %s -o %t --mlir-print-ir-after=convert-func-to-llvm 2>&1 || true ) | FileCheck %s +// RUN: ( mkdir -p %T && ptoas --pto-arch=a5 --pto-backend=vpto --cann-output-version=9.0.0 %s -o %t.cann900 --mlir-print-ir-after=convert-func-to-llvm 2>&1 || true ) | FileCheck %s module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { func.func @simt_float_math_kernel(%dst_f32: !pto.ptr, %dst_f16: !pto.ptr, %dst_bf16: !pto.ptr) attributes {pto.aicore} { @@ -40,6 +41,8 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind f32 %fmin = pto.fmin %f32_a, %f32_b : f32, f32 -> f32 %fmax = pto.fmax %bf16_a, %bf16_b : bf16, bf16 -> bf16 + %fmin16 = pto.fmin %f16_a, %f16_b : f16, f16 -> f16 + %fmax16 = pto.fmax %f16_a, %f16_b : f16, f16 -> f16 %exp = pto.exp %f32_a : f32 -> f32 %exp16 = pto.exp %f16_a : f16 -> f16 %log = pto.log %f32_c : f32 -> f32 @@ -66,6 +69,8 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, f16 pto.store %rint, %dst_f16[%c3] : !pto.ptr, f16 pto.store %fma16, %dst_f16[%c4] : !pto.ptr, f16 + pto.store %fmin16, %dst_f16[%c5] : !pto.ptr, f16 + pto.store %fmax16, %dst_f16[%c6] : !pto.ptr, f16 pto.store %fmax, %dst_bf16[%c0] : !pto.ptr, bf16 pto.store %ceil, %dst_bf16[%c1] : !pto.ptr, bf16 pto.store %round, %dst_bf16[%c2] : !pto.ptr, bf16 @@ -80,6 +85,8 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind None: [3, 2, 2, 1, 2, 2, 0, 16, 16, 1, 2, 2, 2, 6, 6, 6, -4, 4], dtype=np.int32, ) + values = np.arange(-16, 16, dtype=np.float16) + three = np.float16(3) + golden_v1[32:64] = np.minimum(values, three).view(np.uint16).astype(np.int32) + golden_v1[64:96] = np.maximum(values, three).view(np.uint16).astype(np.int32) v1.tofile(output_dir / "v1.bin") golden_v1.tofile(output_dir / "golden_v1.bin") diff --git a/test/vpto/cases/micro-op/simt/simt-float-convert-core/kernel.pto b/test/vpto/cases/micro-op/simt/simt-float-convert-core/kernel.pto index 45f467d886..77725c5d84 100644 --- a/test/vpto/cases/micro-op/simt/simt-float-convert-core/kernel.pto +++ b/test/vpto/cases/micro-op/simt/simt-float-convert-core/kernel.pto @@ -36,6 +36,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind f16 %abs = pto.absf %f32_neg : f32 -> f32 %sqrt_f32 = pto.sqrt %f32_four : f32 -> f32 @@ -87,6 +94,12 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind i32 %s32_f32_i = pto.convert %s32_f32 round(z) sat signed : f32 -> i32 %u32_f16_i = pto.convert %u32_f16 round(z) sat unsigned : f16 -> i32 + %fmin16 = pto.fmin %tid_f16, %f16_three : f16, f16 -> f16 + %fmax16 = pto.fmax %tid_f16, %f16_three : f16, f16 -> f16 + %fmin16_bits = llvm.bitcast %fmin16 : f16 to i16 + %fmax16_bits = llvm.bitcast %fmax16 : f16 to i16 + %fmin16_out = arith.extui %fmin16_bits : i16 to i32 + %fmax16_out = arith.extui %fmax16_bits : i16 to i32 pto.store %abs_i, %dst[%c0] : !pto.ptr, i32 pto.store %sqrt_f32_i, %dst[%c1] : !pto.ptr, i32 @@ -108,6 +121,12 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, i32 pto.store %u32_f16_i, %dst[%c17] : !pto.ptr, i32 + %fmin16_offset = arith.addi %tid, %c32_i32 : i32 + %fmax16_offset = arith.addi %tid, %c64_i32 : i32 + %fmin16_idx = arith.index_castui %fmin16_offset : i32 to index + %fmax16_idx = arith.index_castui %fmax16_offset : i32 to index + pto.store %fmin16_out, %dst[%fmin16_idx] : !pto.ptr, i32 + pto.store %fmax16_out, %dst[%fmax16_idx] : !pto.ptr, i32 return } } From fb59d17296ab073f23cf88d084e92b489385172f Mon Sep 17 00:00:00 2001 From: qukelin Date: Wed, 12 Aug 2026 01:58:48 +0800 Subject: [PATCH 094/122] feat(vpto): support scalar f16 absf --- docs/isa/micro-isa/17-simt.md | 7 ++++--- include/PTO/IR/VPTOOps.td | 8 ++++---- lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp | 2 +- lib/PTO/Transforms/VPTOLLVMEmitter.cpp | 2 +- test/lit/vpto/simt_lowlevel_float_math_vpto_llvm.pto | 4 ++++ .../micro-op/simt/simt-float-convert-core/compare.py | 1 + .../cases/micro-op/simt/simt-float-convert-core/golden.py | 1 + .../micro-op/simt/simt-float-convert-core/kernel.pto | 7 +++++++ 8 files changed, 23 insertions(+), 9 deletions(-) diff --git a/docs/isa/micro-isa/17-simt.md b/docs/isa/micro-isa/17-simt.md index 73db6575a1..3ee361a9bb 100644 --- a/docs/isa/micro-isa/17-simt.md +++ b/docs/isa/micro-isa/17-simt.md @@ -791,10 +791,11 @@ else: - **syntax:** `%r = pto.absf %x : T -> T` - **semantics:** Return `abs(x)`. For `vector<2xT>`, absolute value is applied independently to each element. -- **inputs:** `%x` is an `f32` scalar, `vector<2xf16>`, or `vector<2xbf16>`. +- **inputs:** `%x` is an `f16` or `f32` scalar, `vector<2xf16>`, or + `vector<2xbf16>`. - **outputs:** One value with the same type as `%x`. -- **constraints and limitations:** Scalar `f16` and scalar `bf16` are not - accepted by this op; use the packed form only for `vector<2xT>`. +- **constraints and limitations:** Scalar `bf16` is not accepted by this op; + use the packed form for `vector<2xbf16>`. ### `pto.sqrt` diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index bd1ab9d08b..5d0f5f591f 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -898,10 +898,10 @@ class PTO_BinaryFloatScalarOp def PTO_AbsFOp : PTO_SimtOp<"absf", [Pure, AllTypesMatch<["value", "result"]>]> { - let arguments = (ins AnyTypeOf<[F32, PTO_V2F16Type, PTO_V2BF16Type], - "f32, vector<2xf16> or vector<2xbf16>">:$value); - let results = (outs AnyTypeOf<[F32, PTO_V2F16Type, PTO_V2BF16Type], - "f32, vector<2xf16> or vector<2xbf16>">:$result); + let arguments = (ins AnyTypeOf<[F16, F32, PTO_V2F16Type, PTO_V2BF16Type], + "f16, f32, vector<2xf16> or vector<2xbf16>">:$value); + let results = (outs AnyTypeOf<[F16, F32, PTO_V2F16Type, PTO_V2BF16Type], + "f16, f32, vector<2xf16> or vector<2xbf16>">:$result); let assemblyFormat = [{ $value attr-dict `:` type($value) `->` type($result) diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index 9ce5c0c3c5..db5f9a344e 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -2984,7 +2984,7 @@ template <> FailureOr buildUnaryScalarMathCallee(MLIRContext *context, Type valueType) { std::string elem = getLLVMFloatBuiltinFragment(valueType); - if (elem != "f32" && elem != "v2f16" && elem != "v2bf16") + if (elem != "f16" && elem != "f32" && elem != "v2f16" && elem != "v2bf16") return failure(); return StringAttr::get(context, "llvm.fabs." + elem).getValue(); } diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 68356be08f..9011e1b2e1 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -3012,7 +3012,7 @@ template <> FailureOr buildUnaryScalarMathCallee(MLIRContext *context, Type valueType) { std::string elem = getLLVMFloatBuiltinFragment(valueType); - if (elem != "f32" && elem != "v2f16" && elem != "v2bf16") + if (elem != "f16" && elem != "f32" && elem != "v2f16" && elem != "v2bf16") return failure(); return StringAttr::get(context, "llvm.fabs." + elem).getValue(); } diff --git a/test/lit/vpto/simt_lowlevel_float_math_vpto_llvm.pto b/test/lit/vpto/simt_lowlevel_float_math_vpto_llvm.pto index 4d634c5af7..10d732dea5 100644 --- a/test/lit/vpto/simt_lowlevel_float_math_vpto_llvm.pto +++ b/test/lit/vpto/simt_lowlevel_float_math_vpto_llvm.pto @@ -34,11 +34,13 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind f32 + %abs16 = pto.absf %f16_neg : f16 -> f16 %fmin = pto.fmin %f32_a, %f32_b : f32, f32 -> f32 %fmax = pto.fmax %bf16_a, %bf16_b : bf16, bf16 -> bf16 %fmin16 = pto.fmin %f16_a, %f16_b : f16, f16 -> f16 @@ -64,6 +66,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, f32 pto.store %floor, %dst_f32[%c5] : !pto.ptr, f32 pto.store %fma, %dst_f32[%c6] : !pto.ptr, f32 + pto.store %abs16, %dst_f16[%c7] : !pto.ptr, f16 pto.store %exp16, %dst_f16[%c0] : !pto.ptr, f16 pto.store %log16, %dst_f16[%c1] : !pto.ptr, f16 pto.store %pow16, %dst_f16[%c2] : !pto.ptr, f16 @@ -83,6 +86,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind None: three = np.float16(3) golden_v1[32:64] = np.minimum(values, three).view(np.uint16).astype(np.int32) golden_v1[64:96] = np.maximum(values, three).view(np.uint16).astype(np.int32) + golden_v1[96:128] = np.abs(np.arange(-16, 16, dtype=np.float16)).view(np.uint16) v1.tofile(output_dir / "v1.bin") golden_v1.tofile(output_dir / "golden_v1.bin") diff --git a/test/vpto/cases/micro-op/simt/simt-float-convert-core/kernel.pto b/test/vpto/cases/micro-op/simt/simt-float-convert-core/kernel.pto index 77725c5d84..3a0c7c7a4a 100644 --- a/test/vpto/cases/micro-op/simt/simt-float-convert-core/kernel.pto +++ b/test/vpto/cases/micro-op/simt/simt-float-convert-core/kernel.pto @@ -39,6 +39,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind f16 %abs = pto.absf %f32_neg : f32 -> f32 + %abs_f16 = pto.absf %tid_f16 : f16 -> f16 %sqrt_f32 = pto.sqrt %f32_four : f32 -> f32 %sqrt_f16 = pto.sqrt %f16_four : f16 -> f16 %fmin = pto.fmin %f32_one, %f32_two : f32, f32 -> f32 @@ -98,8 +100,10 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind f16 %fmin16_bits = llvm.bitcast %fmin16 : f16 to i16 %fmax16_bits = llvm.bitcast %fmax16 : f16 to i16 + %abs_f16_bits = llvm.bitcast %abs_f16 : f16 to i16 %fmin16_out = arith.extui %fmin16_bits : i16 to i32 %fmax16_out = arith.extui %fmax16_bits : i16 to i32 + %abs_f16_out = arith.extui %abs_f16_bits : i16 to i32 pto.store %abs_i, %dst[%c0] : !pto.ptr, i32 pto.store %sqrt_f32_i, %dst[%c1] : !pto.ptr, i32 @@ -127,6 +131,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, i32 pto.store %fmax16_out, %dst[%fmax16_idx] : !pto.ptr, i32 + %abs_f16_offset = arith.addi %tid, %c96_i32 : i32 + %abs_f16_idx = arith.index_castui %abs_f16_offset : i32 to index + pto.store %abs_f16_out, %dst[%abs_f16_idx] : !pto.ptr, i32 return } } From 1545250613b24483f37e7d58801f6c59554ae7a1 Mon Sep 17 00:00:00 2001 From: mouliangyu <21963576+mouliangyu@users.noreply.github.com> Date: Tue, 11 Aug 2026 22:05:45 +0800 Subject: [PATCH 095/122] perf(vmi): elide neutral reduction combines --- include/PTO/IR/VMIOps.td | 28 ++- lib/PTO/IR/VMI.cpp | 60 ++---- lib/PTO/Transforms/VMILayoutAssignment.cpp | 12 -- .../Transforms/VMILowerUnifiedToLegacy.cpp | 63 +------ lib/PTO/Transforms/VMIToVPTO.cpp | 176 +++++++++--------- test/lit/vmi_new/vmi_to_vpto_reduce_addf.pto | 3 +- .../vmi_new/vmi_to_vpto_reduce_addf_f16.pto | 3 - .../vmi_to_vpto_reduce_addf_multichunk.pto | 7 +- test/lit/vmi_new/vmi_to_vpto_reduce_addi.pto | 3 +- .../vmi_to_vpto_reduce_addi_multichunk.pto | 7 +- .../vmi_to_vpto_reduce_maxf_multichunk.pto | 12 +- test/lit/vmi_new/vmi_to_vpto_reduce_minf.pto | 4 +- .../vmi_to_vpto_reduce_shape_invalid.pto | 6 +- 13 files changed, 141 insertions(+), 243 deletions(-) diff --git a/include/PTO/IR/VMIOps.td b/include/PTO/IR/VMIOps.td index 1f21c55b67..19399252f4 100644 --- a/include/PTO/IR/VMIOps.td +++ b/include/PTO/IR/VMIOps.td @@ -417,64 +417,58 @@ def VMICompressStoreOp : VMI_Op<"compress_store", [DeclareOpInterfaceMethods { - let summary = "VMI masked integer add reduction with a 1-lane vector init"; + let summary = "VMI masked integer add reduction"; let arguments = (ins VMI_VRegTypeConstraint:$source, - VMI_VRegTypeConstraint:$init, VMI_MaskTypeConstraint:$mask); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; - let assemblyFormat = "$source `,` $init `,` $mask attr-dict `:` type($source) `,` type($init) `,` type($mask) `->` type($result)"; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; } def VMIReduceAddFOp : VMI_Op<"reduce_addf"> { let summary = "VMI masked floating-point add reduction with explicit reassociation permission"; let arguments = (ins VMI_VRegTypeConstraint:$source, - VMI_VRegTypeConstraint:$init, VMI_MaskTypeConstraint:$mask, OptionalAttr:$reassoc); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; - let assemblyFormat = "$source `,` $init `,` $mask attr-dict `:` type($source) `,` type($init) `,` type($mask) `->` type($result)"; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; } def VMIReduceMaxFOp : VMI_Op<"reduce_maxf"> { - let summary = "VMI masked floating-point maximum reduction with a 1-lane vector init"; + let summary = "VMI masked floating-point maximum reduction"; let arguments = (ins VMI_VRegTypeConstraint:$source, - VMI_VRegTypeConstraint:$init, VMI_MaskTypeConstraint:$mask); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; - let assemblyFormat = "$source `,` $init `,` $mask attr-dict `:` type($source) `,` type($init) `,` type($mask) `->` type($result)"; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; } def VMIReduceMinFOp : VMI_Op<"reduce_minf"> { - let summary = "VMI masked floating-point minimum reduction with a 1-lane vector init"; + let summary = "VMI masked floating-point minimum reduction"; let arguments = (ins VMI_VRegTypeConstraint:$source, - VMI_VRegTypeConstraint:$init, VMI_MaskTypeConstraint:$mask); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; - let assemblyFormat = "$source `,` $init `,` $mask attr-dict `:` type($source) `,` type($init) `,` type($mask) `->` type($result)"; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; } def VMIReduceMaxIOp : VMI_Op<"reduce_maxi"> { - let summary = "VMI masked integer maximum reduction with a 1-lane vector init"; + let summary = "VMI masked integer maximum reduction"; let arguments = (ins VMI_VRegTypeConstraint:$source, - VMI_VRegTypeConstraint:$init, VMI_MaskTypeConstraint:$mask); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; - let assemblyFormat = "$source `,` $init `,` $mask attr-dict `:` type($source) `,` type($init) `,` type($mask) `->` type($result)"; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; } def VMIReduceMinIOp : VMI_Op<"reduce_mini"> { - let summary = "VMI masked integer minimum reduction with a 1-lane vector init"; + let summary = "VMI masked integer minimum reduction"; let arguments = (ins VMI_VRegTypeConstraint:$source, - VMI_VRegTypeConstraint:$init, VMI_MaskTypeConstraint:$mask); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; - let assemblyFormat = "$source `,` $init `,` $mask attr-dict `:` type($source) `,` type($init) `,` type($mask) `->` type($result)"; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; } def VMIGroupReduceAddFOp : VMI_Op<"group_reduce_addf"> { diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp index c10315d918..6f5143959e 100644 --- a/lib/PTO/IR/VMI.cpp +++ b/lib/PTO/IR/VMI.cpp @@ -1499,7 +1499,6 @@ void VMICompressStoreOp::getEffects( LogicalResult VMIReduceAddIOp::verify() { auto sourceType = cast(getSource().getType()); - auto initType = cast(getInit().getType()); auto maskType = cast(getMask().getType()); auto resultType = cast(getResult().getType()); if (!isVMIIntegerLikeType(sourceType.getElementType())) @@ -1507,22 +1506,15 @@ LogicalResult VMIReduceAddIOp::verify() { auto sourceIntegerType = dyn_cast(sourceType.getElementType()); if (!sourceIntegerType || sourceIntegerType.getWidth() != 32) return emitOpError("requires 32-bit integer source element type"); - if (sourceType.getElementType() != initType.getElementType() || - sourceType.getElementType() != resultType.getElementType()) - return emitOpError( - "requires source, init, and result element types to match"); - if (initType.getElementCount() != 1 || resultType.getElementCount() != 1) - return emitOpError("requires init and result to be 1-lane VMI vectors"); - if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), - {initType, resultType}, - /*requireSameElement=*/true))) - return failure(); + if (sourceType.getElementType() != resultType.getElementType()) + return emitOpError("requires source and result element types to match"); + if (resultType.getElementCount() != 1) + return emitOpError("requires result to be a 1-lane VMI vector"); return verifyMaskMatchesData(getOperation(), maskType, sourceType); } LogicalResult VMIReduceAddFOp::verify() { auto sourceType = cast(getSource().getType()); - auto initType = cast(getInit().getType()); auto maskType = cast(getMask().getType()); auto resultType = cast(getResult().getType()); if (!getOperation()->hasAttr("reassoc")) @@ -1533,22 +1525,15 @@ LogicalResult VMIReduceAddFOp::verify() { return emitOpError("requires floating-point-like VMI source element type"); if (!isVMIF16OrF32Type(sourceType.getElementType())) return emitOpError("requires f16 or f32 source element type"); - if (sourceType.getElementType() != initType.getElementType() || - sourceType.getElementType() != resultType.getElementType()) - return emitOpError( - "requires source, init, and result element types to match"); - if (initType.getElementCount() != 1 || resultType.getElementCount() != 1) - return emitOpError("requires init and result to be 1-lane VMI vectors"); - if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), - {initType, resultType}, - /*requireSameElement=*/true))) - return failure(); + if (sourceType.getElementType() != resultType.getElementType()) + return emitOpError("requires source and result element types to match"); + if (resultType.getElementCount() != 1) + return emitOpError("requires result to be a 1-lane VMI vector"); return verifyMaskMatchesData(getOperation(), maskType, sourceType); } template LogicalResult verifyReduceMinMaxFOp(OpTy op) { auto sourceType = cast(op.getSource().getType()); - auto initType = cast(op.getInit().getType()); auto maskType = cast(op.getMask().getType()); auto resultType = cast(op.getResult().getType()); if (!isVMIFloatLikeType(sourceType.getElementType())) @@ -1556,16 +1541,10 @@ template LogicalResult verifyReduceMinMaxFOp(OpTy op) { "requires floating-point-like VMI source element type"); if (!isVMIF16OrF32Type(sourceType.getElementType())) return op.emitOpError("requires f16 or f32 source element type"); - if (sourceType.getElementType() != initType.getElementType() || - sourceType.getElementType() != resultType.getElementType()) - return op.emitOpError( - "requires source, init, and result element types to match"); - if (initType.getElementCount() != 1 || resultType.getElementCount() != 1) - return op.emitOpError("requires init and result to be 1-lane VMI vectors"); - if (failed(verifyAllSameVRegShapeAndLayout(op.getOperation(), - {initType, resultType}, - /*requireSameElement=*/true))) - return failure(); + if (sourceType.getElementType() != resultType.getElementType()) + return op.emitOpError("requires source and result element types to match"); + if (resultType.getElementCount() != 1) + return op.emitOpError("requires result to be a 1-lane VMI vector"); return verifyMaskMatchesData(op.getOperation(), maskType, sourceType); } @@ -1575,7 +1554,6 @@ LogicalResult VMIReduceMinFOp::verify() { return verifyReduceMinMaxFOp(*this); } template LogicalResult verifyReduceMinMaxIOp(OpTy op) { auto sourceType = cast(op.getSource().getType()); - auto initType = cast(op.getInit().getType()); auto maskType = cast(op.getMask().getType()); auto resultType = cast(op.getResult().getType()); auto sourceIntegerType = dyn_cast(sourceType.getElementType()); @@ -1583,16 +1561,10 @@ template LogicalResult verifyReduceMinMaxIOp(OpTy op) { !isVMIAnyI8I16I32Type(sourceType.getElementType())) return op.emitOpError( "requires 8-bit, 16-bit, or 32-bit integer source element type"); - if (sourceType.getElementType() != initType.getElementType() || - sourceType.getElementType() != resultType.getElementType()) - return op.emitOpError( - "requires source, init, and result element types to match"); - if (initType.getElementCount() != 1 || resultType.getElementCount() != 1) - return op.emitOpError("requires init and result to be 1-lane VMI vectors"); - if (failed(verifyAllSameVRegShapeAndLayout(op.getOperation(), - {initType, resultType}, - /*requireSameElement=*/true))) - return failure(); + if (sourceType.getElementType() != resultType.getElementType()) + return op.emitOpError("requires source and result element types to match"); + if (resultType.getElementCount() != 1) + return op.emitOpError("requires result to be a 1-lane VMI vector"); return verifyMaskMatchesData(op.getOperation(), maskType, sourceType); } diff --git a/lib/PTO/Transforms/VMILayoutAssignment.cpp b/lib/PTO/Transforms/VMILayoutAssignment.cpp index 49555602fc..47297c389c 100644 --- a/lib/PTO/Transforms/VMILayoutAssignment.cpp +++ b/lib/PTO/Transforms/VMILayoutAssignment.cpp @@ -878,8 +878,6 @@ struct LayoutSolver { if (auto reduce = dyn_cast(op)) { requestDataUse(reduce.getSourceMutable(), getContiguousLayout(), /*late=*/false, DataLayoutSeedPhase::Reduce); - requestDataUse(reduce.getInitMutable(), getContiguousLayout(), - /*late=*/false, DataLayoutSeedPhase::Reduce); if (failed(requestMaskUse(reduce.getMaskMutable(), getContiguousLayout(), op))) return WalkResult::interrupt(); @@ -891,8 +889,6 @@ struct LayoutSolver { if (auto reduce = dyn_cast(op)) { requestDataUse(reduce.getSourceMutable(), getContiguousLayout(), /*late=*/false, DataLayoutSeedPhase::Reduce); - requestDataUse(reduce.getInitMutable(), getContiguousLayout(), - /*late=*/false, DataLayoutSeedPhase::Reduce); if (failed(requestMaskUse(reduce.getMaskMutable(), getContiguousLayout(), op))) return WalkResult::interrupt(); @@ -904,8 +900,6 @@ struct LayoutSolver { if (auto reduce = dyn_cast(op)) { requestDataUse(reduce.getSourceMutable(), getContiguousLayout(), /*late=*/false, DataLayoutSeedPhase::Reduce); - requestDataUse(reduce.getInitMutable(), getContiguousLayout(), - /*late=*/false, DataLayoutSeedPhase::Reduce); if (failed(requestMaskUse(reduce.getMaskMutable(), getContiguousLayout(), op))) return WalkResult::interrupt(); @@ -917,8 +911,6 @@ struct LayoutSolver { if (auto reduce = dyn_cast(op)) { requestDataUse(reduce.getSourceMutable(), getContiguousLayout(), /*late=*/false, DataLayoutSeedPhase::Reduce); - requestDataUse(reduce.getInitMutable(), getContiguousLayout(), - /*late=*/false, DataLayoutSeedPhase::Reduce); if (failed(requestMaskUse(reduce.getMaskMutable(), getContiguousLayout(), op))) return WalkResult::interrupt(); @@ -930,8 +922,6 @@ struct LayoutSolver { if (auto reduce = dyn_cast(op)) { requestDataUse(reduce.getSourceMutable(), getContiguousLayout(), /*late=*/false, DataLayoutSeedPhase::Reduce); - requestDataUse(reduce.getInitMutable(), getContiguousLayout(), - /*late=*/false, DataLayoutSeedPhase::Reduce); if (failed(requestMaskUse(reduce.getMaskMutable(), getContiguousLayout(), op))) return WalkResult::interrupt(); @@ -943,8 +933,6 @@ struct LayoutSolver { if (auto reduce = dyn_cast(op)) { requestDataUse(reduce.getSourceMutable(), getContiguousLayout(), /*late=*/false, DataLayoutSeedPhase::Reduce); - requestDataUse(reduce.getInitMutable(), getContiguousLayout(), - /*late=*/false, DataLayoutSeedPhase::Reduce); if (failed(requestMaskUse(reduce.getMaskMutable(), getContiguousLayout(), op))) return WalkResult::interrupt(); diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp index c4d8066542..41fb5b3285 100644 --- a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp +++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp @@ -144,48 +144,6 @@ static Value createZeroConstant(OpBuilder &builder, Location loc, } -/// Create a 1-lane VMIConstantOp with the neutral element for reduction: -/// add: 0 (int and float) -/// max: -INF (float), INT_MIN (int) -/// min: +INF (float), INT_MAX (int) -static Value createReduceNeutralInit(OpBuilder &builder, Location loc, - Type elemType, bool isAdd, bool isMax, - Attribute layout = Attribute()) { - auto oneLaneType = - VMIVRegType::get(builder.getContext(), 1, elemType, layout); - auto shapedType = RankedTensorType::get({1}, elemType); - DenseElementsAttr attr; - if (auto floatTy = dyn_cast(elemType)) { - if (isAdd) - attr = DenseElementsAttr::get( - shapedType, APFloat::getZero(floatTy.getFloatSemantics())); - else if (isMax) - attr = DenseElementsAttr::get( - shapedType, - APFloat::getInf(floatTy.getFloatSemantics(), /*Negative=*/true)); - else - attr = DenseElementsAttr::get( - shapedType, - APFloat::getInf(floatTy.getFloatSemantics(), /*Negative=*/false)); - } else { - auto intTy = cast(elemType); - if (isAdd) - attr = DenseElementsAttr::get(shapedType, - APInt::getZero(intTy.getWidth())); - else if (isMax) - attr = DenseElementsAttr::get( - shapedType, intTy.isUnsigned() - ? APInt::getZero(intTy.getWidth()) - : APInt::getSignedMinValue(intTy.getWidth())); - else - attr = DenseElementsAttr::get( - shapedType, intTy.isUnsigned() - ? APInt::getMaxValue(intTy.getWidth()) - : APInt::getSignedMaxValue(intTy.getWidth())); - } - return builder.create(loc, oneLaneType, attr).getResult(); -} - /// Map a unified vcmp `cmp` mode to the predicate string for legacy /// cmpf/cmpi. Float operands use ordered predicates (olt, oeq, ...); /// integer operands select signedness from the element type. @@ -826,20 +784,17 @@ static LogicalResult lowerVCadd(VMIvcaddOp op, OpBuilder &builder) { op.getResult().replaceAllUsesWith(result); } else { // Full reduce path - Value init = createReduceNeutralInit(builder, loc, elemType, - /*isAdd=*/true, /*isMax=*/false, - sourceType.getLayout()); Value result; if (isFloat) result = builder - .create(loc, resultType, source, init, mask, + .create(loc, resultType, source, mask, op.getReassocAttr()) .getResult(); else result = builder - .create(loc, resultType, source, init, mask) + .create(loc, resultType, source, mask) .getResult(); op.getResult().replaceAllUsesWith(result); } @@ -880,17 +835,14 @@ static LogicalResult lowerVcmax(VMIvcmaxOp op, OpBuilder &builder) { return success(); } - Value init = createReduceNeutralInit(builder, loc, elemType, - /*isAdd=*/false, /*isMax=*/true, - sourceType.getLayout()); Value result; if (isFloat) result = builder - .create(loc, resultType, source, init, mask) + .create(loc, resultType, source, mask) .getResult(); else result = builder - .create(loc, resultType, source, init, mask) + .create(loc, resultType, source, mask) .getResult(); op.getResult().replaceAllUsesWith(result); op->erase(); @@ -929,17 +881,14 @@ static LogicalResult lowerVcmin(VMIvcminOp op, OpBuilder &builder) { return success(); } - Value init = createReduceNeutralInit(builder, loc, elemType, - /*isAdd=*/false, /*isMax=*/false, - sourceType.getLayout()); Value result; if (isFloat) result = builder - .create(loc, resultType, source, init, mask) + .create(loc, resultType, source, mask) .getResult(); else result = builder - .create(loc, resultType, source, init, mask) + .create(loc, resultType, source, mask) .getResult(); op.getResult().replaceAllUsesWith(result); op->erase(); diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 043329ce59..40a3812266 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -9519,7 +9519,6 @@ struct OneToNVMIReduceAddIOpPattern matchAndRewrite(VMIReduceAddIOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { ValueRange sourceParts = adaptor.getSource(); - ValueRange initParts = adaptor.getInit(); ValueRange maskParts = adaptor.getMask(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); @@ -9527,17 +9526,17 @@ struct OneToNVMIReduceAddIOpPattern return failure(); SmallVector resultTypes = std::move(*maybe_resultTypes); if (sourceParts.empty() || sourceParts.size() != maskParts.size() || - initParts.size() != 1 || resultTypes.size() != 1) + resultTypes.size() != 1) return rewriter.notifyMatchFailure( - op, "reduce_addi requires matching source/mask chunks and one " - "init/result chunk"); + op, "reduce_addi requires matching source/mask chunks and one result " + "chunk"); auto resultType = dyn_cast(resultTypes.front()); auto maskType = dyn_cast(maskParts.front().getType()); - if (!resultType || !maskType || initParts.front().getType() != resultType) + if (!resultType || !maskType) return rewriter.notifyMatchFailure( - op, "reduce_addi requires matching physical source/init/result " - "vregs and one mask"); + op, "reduce_addi requires matching physical source/result vregs and " + "one mask"); for (Value sourcePart : sourceParts) if (sourcePart.getType() != resultType) @@ -9550,12 +9549,6 @@ struct OneToNVMIReduceAddIOpPattern op, "reduce_addi requires every mask chunk to have the same " "predicate type"); - FailureOr firstLaneMask = - createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); - if (failed(firstLaneMask)) - return rewriter.notifyMatchFailure( - op, "failed to create reduce_addi first-lane mask"); - FailureOr combined = combineEquivalentMaskedParts( op.getLoc(), sourceParts, maskParts, resultType, rewriter); if (succeeded(combined)) { @@ -9564,23 +9557,33 @@ struct OneToNVMIReduceAddIOpPattern .create(op.getLoc(), resultType, *combined, maskParts.front()) .getResult(); - Value result = - rewriter - .create(op.getLoc(), resultType, reduced, - initParts.front(), *firstLaneMask) - .getResult(); replaceOpWithFlatConvertedValues( - rewriter, op, SmallVector{result}, + rewriter, op, SmallVector{reduced}, *this->getTypeConverter()); return success(); } - Value accumulator = initParts.front(); - for (auto [sourcePart, maskPart] : - llvm::zip_equal(sourceParts, maskParts)) { + Value accumulator = rewriter + .create(op.getLoc(), resultType, + sourceParts.front(), + maskParts.front()) + .getResult(); + if (sourceParts.size() == 1) { + replaceOpWithFlatConvertedValues( + rewriter, op, SmallVector{accumulator}, + *this->getTypeConverter()); + return success(); + } + FailureOr firstLaneMask = + createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create reduce_addi first-lane mask"); + for (size_t part = 1; part < sourceParts.size(); ++part) { Value reduced = rewriter - .create(op.getLoc(), resultType, sourcePart, maskPart) + .create(op.getLoc(), resultType, sourceParts[part], + maskParts[part]) .getResult(); accumulator = rewriter .create(op.getLoc(), resultType, reduced, @@ -9603,7 +9606,6 @@ struct OneToNVMIReduceAddFOpPattern matchAndRewrite(VMIReduceAddFOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { ValueRange sourceParts = adaptor.getSource(); - ValueRange initParts = adaptor.getInit(); ValueRange maskParts = adaptor.getMask(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); @@ -9611,17 +9613,17 @@ struct OneToNVMIReduceAddFOpPattern return failure(); SmallVector resultTypes = std::move(*maybe_resultTypes); if (sourceParts.empty() || sourceParts.size() != maskParts.size() || - initParts.size() != 1 || resultTypes.size() != 1) + resultTypes.size() != 1) return rewriter.notifyMatchFailure( - op, "reduce_addf requires matching source/mask chunks and one " - "init/result chunk"); + op, "reduce_addf requires matching source/mask chunks and one result " + "chunk"); auto resultType = dyn_cast(resultTypes.front()); auto maskType = dyn_cast(maskParts.front().getType()); - if (!resultType || !maskType || initParts.front().getType() != resultType) + if (!resultType || !maskType) return rewriter.notifyMatchFailure( - op, "reduce_addf requires matching physical source/init/result " - "vregs and one mask"); + op, "reduce_addf requires matching physical source/result vregs and " + "one mask"); for (Value sourcePart : sourceParts) if (sourcePart.getType() != resultType) @@ -9634,12 +9636,6 @@ struct OneToNVMIReduceAddFOpPattern op, "reduce_addf requires every mask chunk to have the same " "predicate type"); - FailureOr firstLaneMask = - createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); - if (failed(firstLaneMask)) - return rewriter.notifyMatchFailure( - op, "failed to create reduce_addf first-lane mask"); - FailureOr combined = combineEquivalentMaskedParts( op.getLoc(), sourceParts, maskParts, resultType, rewriter); if (succeeded(combined)) { @@ -9648,23 +9644,33 @@ struct OneToNVMIReduceAddFOpPattern .create(op.getLoc(), resultType, *combined, maskParts.front()) .getResult(); - Value result = - rewriter - .create(op.getLoc(), resultType, reduced, - initParts.front(), *firstLaneMask) - .getResult(); replaceOpWithFlatConvertedValues( - rewriter, op, SmallVector{result}, + rewriter, op, SmallVector{reduced}, *this->getTypeConverter()); return success(); } - Value accumulator = initParts.front(); - for (auto [sourcePart, maskPart] : - llvm::zip_equal(sourceParts, maskParts)) { + Value accumulator = rewriter + .create(op.getLoc(), resultType, + sourceParts.front(), + maskParts.front()) + .getResult(); + if (sourceParts.size() == 1) { + replaceOpWithFlatConvertedValues( + rewriter, op, SmallVector{accumulator}, + *this->getTypeConverter()); + return success(); + } + FailureOr firstLaneMask = + createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create reduce_addf first-lane mask"); + for (size_t part = 1; part < sourceParts.size(); ++part) { Value reduced = rewriter - .create(op.getLoc(), resultType, sourcePart, maskPart) + .create(op.getLoc(), resultType, sourceParts[part], + maskParts[part]) .getResult(); accumulator = rewriter .create(op.getLoc(), resultType, reduced, @@ -10388,7 +10394,6 @@ struct OneToNVMIReduceMinMaxOpPattern : OpConversionPattern { typename OpConversionPattern::OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { ValueRange sourceParts = adaptor.getSource(); - ValueRange initParts = adaptor.getInit(); ValueRange maskParts = adaptor.getMask(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); @@ -10396,17 +10401,17 @@ struct OneToNVMIReduceMinMaxOpPattern : OpConversionPattern { return failure(); SmallVector resultTypes = std::move(*maybe_resultTypes); if (sourceParts.empty() || sourceParts.size() != maskParts.size() || - initParts.size() != 1 || resultTypes.size() != 1) + resultTypes.size() != 1) return rewriter.notifyMatchFailure( op, "min/max reduction requires matching source/mask chunks " - "and one init/result chunk"); + "and one result chunk"); auto resultType = dyn_cast(resultTypes.front()); auto maskType = dyn_cast(maskParts.front().getType()); - if (!resultType || !maskType || initParts.front().getType() != resultType) + if (!resultType || !maskType) return rewriter.notifyMatchFailure( - op, "min/max reduction requires matching physical source/" - "init/result vregs and one mask"); + op, "min/max reduction requires matching physical source/result " + "vregs and one mask"); for (Value sourcePart : sourceParts) if (sourcePart.getType() != resultType) @@ -10419,12 +10424,6 @@ struct OneToNVMIReduceMinMaxOpPattern : OpConversionPattern { op, "min/max reduction requires every mask chunk to have " "the same predicate type"); - FailureOr firstLaneMask = - createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); - if (failed(firstLaneMask)) - return rewriter.notifyMatchFailure( - op, "failed to create min/max reduction first-lane mask"); - FailureOr combined = combineEquivalentMaskedParts( op.getLoc(), sourceParts, maskParts, resultType, rewriter); if (succeeded(combined)) { @@ -10433,22 +10432,31 @@ struct OneToNVMIReduceMinMaxOpPattern : OpConversionPattern { .create(op.getLoc(), resultType, *combined, maskParts.front()) .getResult(); - Value result = - rewriter - .create(op.getLoc(), resultType, reduced, - initParts.front(), *firstLaneMask) - .getResult(); replaceOpWithFlatConvertedValues( - rewriter, op, SmallVector{result}, *this->getTypeConverter()); + rewriter, op, SmallVector{reduced}, *this->getTypeConverter()); return success(); } - Value accumulator = initParts.front(); - for (auto [sourcePart, maskPart] : - llvm::zip_equal(sourceParts, maskParts)) { + Value accumulator = rewriter + .create(op.getLoc(), resultType, + sourceParts.front(), + maskParts.front()) + .getResult(); + if (sourceParts.size() == 1) { + replaceOpWithFlatConvertedValues( + rewriter, op, SmallVector{accumulator}, + *this->getTypeConverter()); + return success(); + } + FailureOr firstLaneMask = + createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create min/max reduction first-lane mask"); + for (size_t part = 1; part < sourceParts.size(); ++part) { Value reduced = rewriter .create(op.getLoc(), resultType, - sourcePart, maskPart) + sourceParts[part], maskParts[part]) .getResult(); accumulator = rewriter .create(op.getLoc(), resultType, reduced, @@ -12775,18 +12783,16 @@ checkSupportedReduceShape(OpTy op, bool requiresReassoc, return fail("requires reassoc attr for pair-wise floating-point vcadd"); auto sourceType = cast(op.getSource().getType()); - auto initType = cast(op.getInit().getType()); auto maskType = cast(op.getMask().getType()); auto resultType = cast(op.getResult().getType()); VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); - VMILayoutAttr initLayout = initType.getLayoutAttr(); VMILayoutAttr maskLayout = maskType.getLayoutAttr(); VMILayoutAttr resultLayout = resultType.getLayoutAttr(); - if (!sourceLayout || !initLayout || !maskLayout || !resultLayout) - return fail("requires assigned source, init, mask, and result layouts"); - if (!sourceLayout.isContiguous() || !initLayout.isContiguous() || - !maskLayout.isContiguous() || !resultLayout.isContiguous()) - return fail("requires contiguous source, init, mask, and result layouts"); + if (!sourceLayout || !maskLayout || !resultLayout) + return fail("requires assigned source, mask, and result layouts"); + if (!sourceLayout.isContiguous() || !maskLayout.isContiguous() || + !resultLayout.isContiguous()) + return fail("requires contiguous source, mask, and result layouts"); std::string fullChunkReason; if (failed(checkFullDataPhysicalChunks(sourceType, &fullChunkReason))) @@ -12795,17 +12801,15 @@ checkSupportedReduceShape(OpTy op, bool requiresReassoc, fullChunkReason); FailureOr sourceArity = getVMIPhysicalArity(sourceType); - FailureOr initArity = getVMIPhysicalArity(initType); FailureOr maskArity = getVMIPhysicalArity(maskType); FailureOr resultArity = getVMIPhysicalArity(resultType); - if (failed(sourceArity) || failed(initArity) || failed(maskArity) || - failed(resultArity)) + if (failed(sourceArity) || failed(maskArity) || failed(resultArity)) return fail("requires computable physical arity"); if (*sourceArity < 1 || *maskArity != *sourceArity) return fail("requires source and mask physical arity to match and be " "non-empty"); - if (*initArity != 1 || *resultArity != 1) - return fail("requires one init and result physical chunk"); + if (*resultArity != 1) + return fail("requires one result physical chunk"); return success(); } @@ -13697,7 +13701,7 @@ verifySupportedVMIToVPTOOps(ModuleOp module, << kVMIDiagUnsupportedPrefix << "pto.vmi.reduce_addi lowers through pto.vcadd only for " "contiguous full 32-bit integer source chunks with matching " - "mask chunks and one init/result chunk (" + "mask chunks and one result chunk (" << reason << ")"; return WalkResult::interrupt(); } @@ -13711,7 +13715,7 @@ verifySupportedVMIToVPTOOps(ModuleOp module, << kVMIDiagUnsupportedPrefix << "pto.vmi.reduce_addf lowers through pto.vcadd only with " "reassoc, f32 contiguous full source chunks, matching mask " - "chunks, and one init/result chunk (" + "chunks, and one result chunk (" << reason << ")"; return WalkResult::interrupt(); } @@ -13812,7 +13816,7 @@ verifySupportedVMIToVPTOOps(ModuleOp module, << kVMIDiagUnsupportedPrefix << "pto.vmi.reduce_maxf lowers through pto.vcmax only for f16/f32 " "contiguous full source chunks with matching mask chunks and one " - "init/result chunk (" + "result chunk (" << reason << ")"; return WalkResult::interrupt(); } @@ -13826,7 +13830,7 @@ verifySupportedVMIToVPTOOps(ModuleOp module, << kVMIDiagUnsupportedPrefix << "pto.vmi.reduce_minf lowers through pto.vcmin only for f16/f32 " "contiguous full source chunks with matching mask chunks and one " - "init/result chunk (" + "result chunk (" << reason << ")"; return WalkResult::interrupt(); } @@ -13840,7 +13844,7 @@ verifySupportedVMIToVPTOOps(ModuleOp module, << kVMIDiagUnsupportedPrefix << "pto.vmi.reduce_maxi lowers through pto.vcmax only for " "contiguous full integer source chunks with matching mask " - "chunks and one init/result chunk (" + "chunks and one result chunk (" << reason << ")"; return WalkResult::interrupt(); } @@ -13854,7 +13858,7 @@ verifySupportedVMIToVPTOOps(ModuleOp module, << kVMIDiagUnsupportedPrefix << "pto.vmi.reduce_mini lowers through pto.vcmin only for " "contiguous full integer source chunks with matching mask " - "chunks and one init/result chunk (" + "chunks and one result chunk (" << reason << ")"; return WalkResult::interrupt(); } diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_addf.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_addf.pto index b347a05709..17c72f287f 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_addf.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_addf.pto @@ -25,9 +25,8 @@ module { } // CHECK-LABEL: func.func @vmi_to_vpto_reduce_addf( -// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[REDUCED:.*]] = pto.vcadd %arg0, %arg1 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: pto.vadd %[[REDUCED]], {{.*}}, %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: return %[[REDUCED]] : !pto.vreg<64xf32> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_addf_f16.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_addf_f16.pto index 4ab1698bb9..5c14eb4b93 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_addf_f16.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_addf_f16.pto @@ -25,11 +25,8 @@ module { } // CHECK-LABEL: func.func @vmi_to_vpto_reduce_addf_f16( -// CHECK: %[[LANE0:.*]] = pto.pset_b16 "PAT_VL1" : !pto.mask // CHECK: %[[REDUCED:.*]] = pto.vcadd %arg0, %arg1 // CHECK-SAME: !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<128xf16> -// CHECK: pto.vadd %[[REDUCED]], {{.*}}, %[[LANE0]] -// CHECK-SAME: !pto.vreg<128xf16>, !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<128xf16> // CHECK: return {{.*}} : !pto.vreg<128xf16> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_addf_multichunk.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_addf_multichunk.pto index b0e9f7e7ec..eeed5592d0 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_addf_multichunk.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_addf_multichunk.pto @@ -41,11 +41,10 @@ module { } // CHECK-LABEL: func.func @vmi_to_vpto_reduce_addf_multichunk( -// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[RED0:.*]] = pto.vcadd %arg0, %arg2 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: pto.vadd %[[RED0]], {{.*}}, %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[RED1:.*]] = pto.vcadd %arg1, %arg3 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: pto.vadd %[[RED1]], {{.*}}, %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: pto.vadd %[[RED1]], %[[RED0]], %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast @@ -55,7 +54,7 @@ module { // CHECK: pto.pset_b32 "PAT_ALL" : !pto.mask // CHECK: %[[MERGED:.*]] = pto.vadd %arg0, %arg1, %[[MASK0]] // CHECK: %[[REDUCED:.*]] = pto.vcadd %[[MERGED]], %[[MASK0]] -// CHECK: pto.vadd %[[REDUCED]], {{.*}} +// CHECK-NOT: pto.vadd %[[REDUCED]], // CHECK-NOT: pto.vcadd // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_addi.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_addi.pto index 40d1075b73..dfbab711d5 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_addi.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_addi.pto @@ -25,9 +25,8 @@ module { } // CHECK-LABEL: func.func @vmi_to_vpto_reduce_addi( -// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[REDUCED:.*]] = pto.vcadd %arg0, %arg1 : !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> -// CHECK: pto.vadd %[[REDUCED]], {{.*}}, %[[FIRST]] : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> +// CHECK: return %[[REDUCED]] : !pto.vreg<64xi32> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_addi_multichunk.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_addi_multichunk.pto index 8972d51d90..59cb27f6e2 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_addi_multichunk.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_addi_multichunk.pto @@ -41,11 +41,10 @@ module { } // CHECK-LABEL: func.func @vmi_to_vpto_reduce_addi_multichunk( -// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[RED0:.*]] = pto.vcadd %arg0, %arg2 : !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> -// CHECK: pto.vadd %[[RED0]], {{.*}}, %[[FIRST]] : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> +// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[RED1:.*]] = pto.vcadd %arg1, %arg3 : !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> -// CHECK: pto.vadd %[[RED1]], {{.*}}, %[[FIRST]] : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> +// CHECK: pto.vadd %[[RED1]], %[[RED0]], %[[FIRST]] : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast @@ -55,7 +54,7 @@ module { // CHECK: pto.pset_b32 "PAT_ALL" : !pto.mask // CHECK: %[[MERGED:.*]] = pto.vadd %arg0, %arg1, %[[MASK0]] // CHECK: %[[REDUCED:.*]] = pto.vcadd %[[MERGED]], %[[MASK0]] -// CHECK: pto.vadd %[[REDUCED]], {{.*}} +// CHECK-NOT: pto.vadd %[[REDUCED]], // CHECK-NOT: pto.vcadd // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_maxf_multichunk.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_maxf_multichunk.pto index 95238ab820..8ebe3ff3a3 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_maxf_multichunk.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_maxf_multichunk.pto @@ -55,21 +55,19 @@ module { } // CHECK-LABEL: func.func @vmi_to_vpto_reduce_maxf_multichunk( -// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[RED0:.*]] = pto.vcmax %arg0, %arg2 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: pto.vmax %[[RED0]], {{.*}}, %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[RED1:.*]] = pto.vcmax %arg1, %arg3 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: pto.vmax %[[RED1]], {{.*}}, %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: pto.vmax %[[RED1]], %[[RED0]], %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast // CHECK-LABEL: func.func @vmi_to_vpto_reduce_minf_multichunk( -// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[RED0:.*]] = pto.vcmin %arg0, %arg2 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: pto.vmin %[[RED0]], {{.*}}, %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[RED1:.*]] = pto.vcmin %arg1, %arg3 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: pto.vmin %[[RED1]], {{.*}}, %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: pto.vmin %[[RED1]], %[[RED0]], %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast @@ -79,7 +77,7 @@ module { // CHECK: pto.pset_b32 "PAT_ALL" : !pto.mask // CHECK: %[[MERGED:.*]] = pto.vmax %arg0, %arg1, %[[MASK0]] // CHECK: %[[REDUCED:.*]] = pto.vcmax %[[MERGED]], %[[MASK0]] -// CHECK: pto.vmax %[[REDUCED]], {{.*}} +// CHECK-NOT: pto.vmax %[[REDUCED]], // CHECK-NOT: pto.vcmax // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_minf.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_minf.pto index 9fd5011aba..d5be772ee9 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_minf.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_minf.pto @@ -25,10 +25,8 @@ module { } // CHECK-LABEL: func.func @vmi_to_vpto_reduce_minf( -// CHECK: %[[FIRST:.*]] = pto.pset_b16 "PAT_VL1" : !pto.mask // CHECK: %[[REDUCED:.*]] = pto.vcmin %arg0, %arg1 : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<128xf16> -// CHECK: %[[OUT:.*]] = pto.vmin %[[REDUCED]], {{.*}}, %[[FIRST]] : !pto.vreg<128xf16>, !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<128xf16> -// CHECK: return %[[OUT]] +// CHECK: return %[[REDUCED]] : !pto.vreg<128xf16> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_shape_invalid.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_shape_invalid.pto index 9db5cf6071..8ce979d812 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_shape_invalid.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_shape_invalid.pto @@ -38,7 +38,8 @@ module { } } -// CHECK: error: 'pto.vmi.reduce_addf' op requires all layout-assigned VMI data values to have the same layout +// CHECK: VMI{{-}}UNSUPPORTED{{:}} pto.vmi.reduce_addf lowers through pto.vcadd only +// CHECK-SAME: requires contiguous source, mask, and result layouts // ----- @@ -72,4 +73,5 @@ module { } } -// CHECK: error: 'pto.vmi.reduce_maxf' op requires all layout-assigned VMI data values to have the same layout +// CHECK: VMI{{-}}UNSUPPORTED{{:}} pto.vmi.reduce_maxf lowers through pto.vcmax only +// CHECK-SAME: requires contiguous source, mask, and result layouts From 1471f342bab3eea82f6dd9947c903055c1517fc6 Mon Sep 17 00:00:00 2001 From: jimmychou <47636600+jimmychou0@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:11:19 +0800 Subject: [PATCH 096/122] fix: complete metrics rule cleanup --- .../GraphSyncSolver/EventIdSolver.h | 5 +- lib/Bindings/Python/PTOModule.cpp | 19 +- lib/CAPI/Dialect/PTO.cpp | 54 +- lib/PTO/IR/PTO.cpp | 218 +- lib/PTO/IR/PTOAttrs.cpp | 113 +- lib/PTO/IR/PTOTypeDefs.cpp | 40 +- lib/PTO/IR/PTOTypeUtils.cpp | 9 +- lib/PTO/IR/VMI.cpp | 1016 +++++--- lib/PTO/IR/VPTO.cpp | 2006 ++++++++++----- .../Transforms/BufidSync/BufidSyncAnalysis.h | 2 - .../Transforms/BufidSync/BufidSyncIdAlloc.cpp | 2 +- .../Transforms/BufidSync/BufidSyncIdAlloc.h | 2 +- lib/PTO/Transforms/ConvertToPTOOp.cpp | 12 +- lib/PTO/Transforms/CppPostprocess.cpp | 27 +- lib/PTO/Transforms/ExpandTileOp.cpp | 298 ++- lib/PTO/Transforms/FoldTileBufIntrinsics.cpp | 153 +- .../Transforms/GraphSyncSolver/Utility.cpp | 1 - lib/PTO/Transforms/InferPTOLayout.cpp | 122 +- lib/PTO/Transforms/InferPTOMemScope.cpp | 107 +- .../InsertSync/MemoryDependentAnalyzer.cpp | 4 +- lib/PTO/Transforms/InsertSync/SyncCommon.cpp | 4 +- .../Transforms/InsertTemplateAttributes.cpp | 180 +- lib/PTO/Transforms/LowerPTOToUBufOps.cpp | 390 ++- lib/PTO/Transforms/PTOA5NormalizeTMovPass.cpp | 21 +- .../PTOAssignDefaultFrontendPipeIdPass.cpp | 3 +- lib/PTO/Transforms/PTOCanonicalizeIR.cpp | 12 +- lib/PTO/Transforms/PTOInferVPTOVecScope.cpp | 234 +- .../PTOInferValidatePipeInitPass.cpp | 61 +- .../PTOInstantiateAndInlineOpLib.cpp | 87 +- .../PTOLowerFrontendPipeOpsPass.cpp | 48 +- lib/PTO/Transforms/PTOLowerToOpLibCalls.cpp | 6 +- .../PTOMaterializeSIMTPersistentFragment.cpp | 23 +- .../PTOMaterializeTileOpSections.cpp | 66 +- .../Transforms/PTONarrowVPTOLoopCounters.cpp | 3 +- .../PTONormalizeUncoveredTileSections.cpp | 248 +- lib/PTO/Transforms/PTOOutlineSIMTSections.cpp | 57 +- lib/PTO/Transforms/PTOPlanMemory.cpp | 301 ++- lib/PTO/Transforms/PTOPlanMemory.h | 7 +- lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 513 ++-- .../PTORematerializeFixpipeVectorQuant.cpp | 6 +- lib/PTO/Transforms/PTORemoveIdentityTMov.cpp | 108 +- .../Transforms/PTORemoveRedundantBarrier.cpp | 105 +- lib/PTO/Transforms/PTOResolveBufferSelect.cpp | 48 +- .../PTOResolveReservedBuffersPass.cpp | 38 +- lib/PTO/Transforms/PTOToEmitC.cpp | 268 +- lib/PTO/Transforms/PTOVPTOPtrBoundary.cpp | 84 +- .../Transforms/PTOValidateIntToPtrUses.cpp | 3 +- .../PTOValidatePhysicalSectionBoundaries.cpp | 30 +- lib/PTO/Transforms/PTOValidateVMIIR.cpp | 267 +- lib/PTO/Transforms/PTOValidateVPTOIR.cpp | 263 +- lib/PTO/Transforms/PTOVerifyTFreePass.cpp | 18 +- .../SIMTPersistentFragmentAnalysis.cpp | 12 +- lib/PTO/Transforms/SlotAffineAnalysis.cpp | 21 +- .../Transforms/TileFusion/FusionAnalysis.cpp | 208 +- .../TileFusion/FusionOpSemantics.cpp | 8 +- .../TileFusion/PTOFlattenFusionRegion.cpp | 6 +- .../TileFusion/PTOFusionLoadStoreElision.cpp | 175 +- .../Transforms/TileFusion/PTOFusionPlan.cpp | 78 +- .../TileFusion/PTOFusionPredicateElision.cpp | 105 +- .../TileFusion/PTOFusionRegionGen.cpp | 126 +- .../TileFusion/PTOLowLevelLoopFusion.cpp | 126 +- .../Transforms/TileFusion/PTOMarkLastUse.cpp | 57 +- .../Transforms/TileFusion/PTOOpScheduling.cpp | 66 +- .../TileFusion/PTOPreFusionAnalysis.cpp | 3 +- .../TileFusion/PTOPrintPreFusionAnalysis.cpp | 16 +- .../TileFusion/PTOUnrollAfterLoopFusion.cpp | 9 +- lib/PTO/Transforms/Utils.cpp | 34 +- lib/PTO/Transforms/VMILayoutAssignment.cpp | 16 +- lib/PTO/Transforms/VMILayoutSupport.cpp | 22 +- .../Transforms/VMILowerUnifiedToLegacy.cpp | 1 - .../VMIMaskGranularityAssignment.cpp | 6 +- .../VMINormalizeSignlessIntToUnsigned.cpp | 2 +- lib/PTO/Transforms/VMIToVPTO.cpp | 2218 +++++++++++------ .../Transforms/VPTOBufferMaterialization.cpp | 25 +- lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp | 175 +- lib/PTO/Transforms/VPTOExpandWrapperOps.cpp | 252 +- lib/PTO/Transforms/VPTOLLVMEmitter.cpp | 1471 ++++++++++- .../Transforms/VPTOLLVMEmitterDispatcher.cpp | 24 +- lib/PTO/Transforms/VPTOLLVMEmitterHelper.cpp | 133 +- lib/PTO/Transforms/VPTOMaskSimplify.cpp | 6 +- lib/PTO/Transforms/VPTONormalizeContainer.cpp | 12 +- lib/PTO/Transforms/VPTOOptimizeVcvt.cpp | 90 +- lib/PTO/Transforms/VPTOPtrCastCleanup.cpp | 24 +- lib/PTO/Transforms/VPTOPtrNormalize.cpp | 246 +- lib/PTO/Transforms/VPTOSoftPostUpdate.cpp | 672 +++-- ptodsl/ptodsl/_ops.py | 2 - tools/ptoas/NativeModule.cpp | 3 +- tools/ptoas/ObjectEmission.cpp | 195 +- tools/ptoas/VFSIMTSizePatcher.cpp | 90 +- tools/ptoas/VPTOHostStubEmission.cpp | 30 +- tools/ptoas/driver.cpp | 225 +- tools/ptoas/ptoas.cpp | 636 +++-- tools/ptobc/src/canonical_printer.cpp | 36 +- tools/ptobc/src/leb128.cpp | 6 +- tools/ptobc/src/mlir_encode.cpp | 2 +- tools/ptobc/src/ptobc_format.cpp | 2 +- 96 files changed, 10691 insertions(+), 4693 deletions(-) diff --git a/include/PTO/Transforms/GraphSyncSolver/EventIdSolver.h b/include/PTO/Transforms/GraphSyncSolver/EventIdSolver.h index 96d933fa95..907a9b42b3 100644 --- a/include/PTO/Transforms/GraphSyncSolver/EventIdSolver.h +++ b/include/PTO/Transforms/GraphSyncSolver/EventIdSolver.h @@ -35,7 +35,7 @@ class Action { public: const ACTION_TYPE actionType; Action() = delete; - Action(ACTION_TYPE actionType) : actionType(actionType) {}; + explicit Action(ACTION_TYPE actionType) : actionType(actionType) {}; virtual ~Action() = default; virtual std::string str() const = 0; }; @@ -52,7 +52,8 @@ class ActionNone : public Action { class ActionAddNode : public Action { public: EventIdNode *const node; - ActionAddNode(EventIdNode *node) : Action(ACTION_TYPE::ADD_NODE), node(node) { + explicit ActionAddNode(EventIdNode *node) + : Action(ACTION_TYPE::ADD_NODE), node(node) { assert(node != nullptr); } static bool classof(const Action *e) { diff --git a/lib/Bindings/Python/PTOModule.cpp b/lib/Bindings/Python/PTOModule.cpp index 94d3f728f4..a80981da98 100644 --- a/lib/Bindings/Python/PTOModule.cpp +++ b/lib/Bindings/Python/PTOModule.cpp @@ -114,8 +114,7 @@ static MlirAttribute optionalAttributeFromPy(py::object attr) { return py::cast(attr); } -void populatePTODialectSubmodule(pybind11::module &m); -void populatePTODialectSubmodule(pybind11::module &m) { +static void populatePTODialectSubmodule(pybind11::module &m) { (void)m; } @@ -1204,7 +1203,7 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { std::vector shp = toShapeVectorOrDynamicRank(shape_or_rank); context = inferContextFromElementType(context, elementType); MlirType t = mlirPTOTensorViewTypeGet( - context, (intptr_t)shp.size(), shp.data(), elementType); + context, static_cast(shp.size()), shp.data(), elementType); return cls.attr("__call__")(t); }, py::arg("cls"), py::arg("shape_or_rank"), py::arg("element_type"), @@ -1236,7 +1235,7 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { std::vector shp = toShapeVectorOrDynamicRank(shape_or_rank); context = inferContextFromElementType(context, elementType); MlirType t = mlirPTOPartitionTensorViewTypeGet(context, - (intptr_t)shp.size(), + static_cast(shp.size()), shp.data(), elementType); return cls.attr("__call__")(t); @@ -1268,7 +1267,7 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { [](py::object cls, py::sequence shape, MlirType elementType, MlirContext context) -> py::object { auto shp = toInt64Vector(shape); MlirType t = mlirPTOTileTypeGet(context, - (intptr_t)shp.size(), + static_cast(shp.size()), shp.data(), elementType); return cls.attr("__call__")(t); @@ -1359,7 +1358,7 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { if (!validShapeObj.is_none()) { // 支持 valid_shape 为 list[int] 或 list[Optional[int]] py::list lst = validShapeObj.cast(); - if ((size_t)lst.size() != shape.size()) { + if (static_cast(lst.size()) != shape.size()) { throw std::runtime_error("valid_shape rank must match shape rank"); } validShape.resize(lst.size()); @@ -1379,16 +1378,16 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { MlirAttribute cfg = configObj.cast(); ty = mlirPTOTileBufTypeGetWithValidShapeAndConfig( ctx, - (intptr_t)shape.size(), shape.data(), + static_cast(shape.size()), shape.data(), elementType, memorySpace, - (intptr_t)validShape.size(), validShape.data(), + static_cast(validShape.size()), validShape.data(), cfg); } else { ty = mlirPTOTileBufTypeGetWithValidShape( ctx, - (intptr_t)shape.size(), shape.data(), + static_cast(shape.size()), shape.data(), elementType, memorySpace, - (intptr_t)validShape.size(), validShape.data()); + static_cast(validShape.size()), validShape.data()); } if (mlirTypeIsNull(ty)) return py::none(); diff --git a/lib/CAPI/Dialect/PTO.cpp b/lib/CAPI/Dialect/PTO.cpp index 40520b90fc..61a37c11fa 100644 --- a/lib/CAPI/Dialect/PTO.cpp +++ b/lib/CAPI/Dialect/PTO.cpp @@ -53,8 +53,9 @@ static CanonicalValidShapeVector canonicalizeTileBufValidShape(ArrayRef validShape) { CanonicalValidShapeVector canonical; canonical.reserve(validShape.size()); - for (int64_t dim : validShape) + for (int64_t dim : validShape) { canonical.push_back(dim < 0 ? ShapedType::kDynamic : dim); + } return canonical; } @@ -298,7 +299,9 @@ MlirType mlirPTOTileBufTypeGetWithConfig(MlirContext ctx, intptr_t rank, MLIRContext *c = unwrap(ctx); auto shp = llvm::ArrayRef(shape, rank); auto cfg = mlir::dyn_cast_or_null(unwrap(config)); - if (!cfg) cfg = mlir::pto::TileBufConfigAttr::getDefault(c); + if (!cfg) { + cfg = mlir::pto::TileBufConfigAttr::getDefault(c); + } auto ty = mlir::pto::TileBufType::get(c, shp, unwrap(elementType), unwrap(memorySpace), cfg); return wrap(ty); } @@ -615,16 +618,18 @@ MlirAttribute mlirPTOMaskPatternAttrGet(MlirContext ctx, int32_t value) { default: break; } - if (!v) + if (!v) { return MlirAttribute{nullptr}; + } return wrap(mlir::pto::MaskPatternAttr::get(c, *v)); } MlirAttribute mlirPTOMaskPatternAttrGetLegacyRaw(MlirContext ctx, int32_t value) { auto *c = unwrap(ctx); std::optional v = maskPatternFromLegacyRaw(value); - if (!v) + if (!v) { return MlirAttribute{nullptr}; + } return wrap(mlir::pto::MaskPatternAttr::get(c, *v)); } @@ -642,8 +647,9 @@ MlirAttribute mlirPTOMaskPatternAttrGetEnum(MlirContext ctx, auto *c = unwrap(ctx); std::optional v = maskPatternFromIsaValue(static_cast(value)); - if (!v) + if (!v) { return MlirAttribute{nullptr}; + } return wrap(mlir::pto::MaskPatternAttr::get(c, *v)); } @@ -692,30 +698,41 @@ MlirAttribute mlirPTOTileBufConfigAttrGetDefault(MlirContext ctx) { } static mlir::pto::BLayoutAttr toBLayoutAttr(mlir::MLIRContext *c, mlir::Attribute a) { - if (auto bl = mlir::dyn_cast(a)) return bl; - if (auto ia = mlir::dyn_cast(a)) + if (auto bl = mlir::dyn_cast(a)) { + return bl; + } + if (auto ia = mlir::dyn_cast(a)) { return mlir::pto::BLayoutAttr::get(c, static_cast(ia.getInt())); + } return {}; } static mlir::pto::SLayoutAttr toSLayoutAttr(mlir::MLIRContext *c, mlir::Attribute a) { - if (auto sl = mlir::dyn_cast(a)) return sl; - if (auto ia = mlir::dyn_cast(a)) + if (auto sl = mlir::dyn_cast(a)) { + return sl; + } + if (auto ia = mlir::dyn_cast(a)) { return mlir::pto::SLayoutAttr::get(c, static_cast(ia.getInt())); + } return {}; } static mlir::pto::PadValueAttr toPadValueAttr(mlir::MLIRContext *c, mlir::Attribute a) { - if (auto pv = mlir::dyn_cast(a)) return pv; - if (auto ia = mlir::dyn_cast(a)) + if (auto pv = mlir::dyn_cast(a)) { + return pv; + } + if (auto ia = mlir::dyn_cast(a)) { return mlir::pto::PadValueAttr::get(c, static_cast(ia.getInt())); + } return {}; } static mlir::pto::CompactModeAttr toCompactModeAttr(mlir::MLIRContext *c, mlir::Attribute a) { - if (auto cm = mlir::dyn_cast(a)) + if (auto cm = mlir::dyn_cast(a)) { return cm; - if (auto ia = mlir::dyn_cast(a)) + } + if (auto ia = mlir::dyn_cast(a)) { return mlir::pto::CompactModeAttr::get( c, static_cast(ia.getInt())); + } return {}; } @@ -859,12 +876,14 @@ MlirAttribute mlirPTOTileBufConfigAttrGetWithCompactMode( auto slA = toSLayoutAttr(c, unwrap(sLayout)); auto pvA = toPadValueAttr(c, unwrap(pad)); auto cmA = toCompactModeAttr(c, unwrap(compactMode)); - if (!blA || !slA || !pvA || !cmA) + if (!blA || !slA || !pvA || !cmA) { return MlirAttribute{nullptr}; + } auto sz = mlir::dyn_cast(unwrap(sFractalSize)); - if (!sz || !sz.getType().isInteger(kI32BitWidth)) + if (!sz || !sz.getType().isInteger(kI32BitWidth)) { return MlirAttribute{nullptr}; + } return wrap(mlir::pto::TileBufConfigAttr::get(c, blA, slA, sz, pvA, cmA)); } @@ -877,8 +896,9 @@ MlirType mlirPTOGMTypeGet(MlirContext ctx, intptr_t rank, const int64_t *shape, llvm::SmallVector strides( static_cast(rank), ShapedType::kDynamic); - if (rank > 0) - strides[static_cast(rank) - 1] = 1; + if (rank > 0) { + strides[static_cast(rank) - 1] = 1; + } auto layout = StridedLayoutAttr::get(c, ShapedType::kDynamic, llvm::ArrayRef(strides)); auto memSpace = mlir::pto::AddressSpaceAttr::get(c, mlir::pto::AddressSpace::GM); diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index aa0a92b816..a33e08861d 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -1212,7 +1212,6 @@ void mlir::pto::TScatterOp::print(OpAsmPrinter &p) { } namespace { - struct CommRecvClause { OpAsmParser::UnresolvedOperand ping; std::optional pong; @@ -1941,7 +1940,7 @@ inferLayout(ArrayRef shape, ArrayRef strides, // ND: row-major contiguous bool isRowMajor = true; - for (int i = 0, e = (int)shape.size() - 1; i < e; ++i) { + for (int i = 0, e = static_cast(shape.size()) - 1; i < e; ++i) { auto expectedStride = multiplyLayoutInts(strides[i + 1], shape[i + 1]); if (!expectedStride || strides[i] != *expectedStride) { isRowMajor = false; @@ -1953,7 +1952,7 @@ inferLayout(ArrayRef shape, ArrayRef strides, // DN: col-major bool isColMajor = true; - for (int i = 0, e = (int)shape.size() - 1; i < e; ++i) { + for (int i = 0, e = static_cast(shape.size()) - 1; i < e; ++i) { auto expectedStride = multiplyLayoutInts(strides[i], shape[i]); if (!expectedStride || strides[i + 1] != *expectedStride) { isColMajor = false; @@ -3958,10 +3957,18 @@ static bool isTileLikeType(Type ty) { } static Type getElemTy(Type ty) { - if (auto tt = mlir::dyn_cast(ty)) return tt.getElementType(); - if (auto tv = mlir::dyn_cast(ty)) return tv.getElementType(); - if (auto tb = mlir::dyn_cast(ty)) return tb.getElementType(); - if (auto tv = mlir::dyn_cast(ty)) return tv.getElementType(); + if (auto tt = mlir::dyn_cast(ty)) { + return tt.getElementType(); + } + if (auto tv = mlir::dyn_cast(ty)) { + return tv.getElementType(); + } + if (auto tb = mlir::dyn_cast(ty)) { + return tb.getElementType(); + } + if (auto tv = mlir::dyn_cast(ty)) { + return tv.getElementType(); + } return Type(); } @@ -4057,8 +4064,9 @@ static LogicalResult verifyAsyncFlatContiguous1DGMViewLike(Operation *op, } bool logical1D = true; - for (int i = 0, e = static_cast(shape.size()) - 1; i < e; ++i) + for (int i = 0, e = static_cast(shape.size()) - 1; i < e; ++i) { logical1D &= shape[i] == 1; + } if (!logical1D) return op->emitOpError() << "expects " << name @@ -5965,7 +5973,6 @@ LogicalResult pto::TCIOp::verify() { } LogicalResult pto::TTriOp::verify() { - Type dstTy = getDst().getType(); if (failed(verifyVecTileCommon(*this, dstTy, "dst"))) return failure(); @@ -6589,7 +6596,6 @@ llvm::LogicalResult mlir::pto::TRandomOp::verify() { return emitOpError("trandom is only supported for A5 targets"); }; auto verifyA5 = [&]() -> LogicalResult { - Type dstTy = getDst().getType(); if (failed(verifyTileBufCommon(*this, dstTy, "dst"))) return failure(); @@ -7115,10 +7121,13 @@ mlir::LogicalResult mlir::pto::TInsertOp::verify() { return emitOpError("fp is only valid with src loc=acc"); auto fpTy = getFp().getType(); auto fpTb = dyn_cast(fpTy); - if (!fpTb) return emitOpError("expects fp to be !pto.tile_buf"); + if (!fpTb) { + return emitOpError("expects fp to be !pto.tile_buf"); + } if (failed(verifyTileBufCommon(*this, fpTy, "fp", - /*allowLowPrecision=*/isA5))) + /*allowLowPrecision=*/isA5))) { return failure(); + } auto fpSpace = getSpace(fpTy); if (!fpSpace || *fpSpace != pto::AddressSpace::SCALING) return emitOpError("expects fp to be loc=scaling"); @@ -8006,9 +8015,15 @@ mlir::LogicalResult mlir::pto::TMovOp::verify() { // 辅助函数：获取 Rank，支持 ShapedType 和 PTO TileTypes static int64_t getRankHelper(Type t) { - if (auto s = dyn_cast(t)) return s.getRank(); - if (auto tile = dyn_cast(t)) return tile.getRank(); - if (auto view = dyn_cast(t)) return view.getRank(); + if (auto s = dyn_cast(t)) { + return s.getRank(); + } + if (auto tile = dyn_cast(t)) { + return tile.getRank(); + } + if (auto view = dyn_cast(t)) { + return view.getRank(); + } return -1; } @@ -9136,7 +9151,6 @@ void mlir::pto::MScatterOp::print(OpAsmPrinter &p) { } LogicalResult MScatterOp::verify() { - Type srcTy = getSrc().getType(); Type idxTy = getIdx().getType(); Type memTy = getMem().getType(); @@ -9363,7 +9377,6 @@ void mlir::pto::MGatherOp::print(OpAsmPrinter &p) { } LogicalResult MGatherOp::verify() { - Type memTy = getMem().getType(); Type idxTy = getIdx().getType(); Type dstTy = getDst().getType(); @@ -10601,7 +10614,6 @@ mlir::LogicalResult mlir::pto::TQuantMxOp::verify() { }; auto verifyA5 = [&]() -> LogicalResult { - Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); Type expTy = getExp().getType(); @@ -10808,7 +10820,6 @@ mlir::LogicalResult mlir::pto::TReluOp::verify() { mlir::LogicalResult mlir::pto::TRemOp::verify() { - Type src0Ty = getSrc0().getType(); Type src1Ty = getSrc1().getType(); Type dstTy = getDst().getType(); @@ -10963,7 +10974,6 @@ mlir::LogicalResult mlir::pto::TRemSOp::verify() { } mlir::LogicalResult mlir::pto::TFModSOp::verify() { - Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); Type scalarTy = getScalar().getType(); @@ -11002,7 +11012,6 @@ static LogicalResult verifyTPowTmpShape(Operation *op, Type tmpTy, Type dstTy) { } mlir::LogicalResult mlir::pto::TPowOp::verify() { - Type baseTy = getBase().getType(); Type expTy = getExp().getType(); Type dstTy = getDst().getType(); @@ -11069,7 +11078,6 @@ mlir::LogicalResult mlir::pto::TPowOp::verify() { } mlir::LogicalResult mlir::pto::TPowSOp::verify() { - Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); Type scalarTy = getScalar().getType(); @@ -12641,7 +12649,9 @@ mlir::LogicalResult mlir::pto::TScatterOp::verify() { } auto isAllowedDataElem = [&](mlir::Type t) -> bool { - if (t.isF16() || t.isF32() || t.isBF16()) return true; + if (t.isF16() || t.isF32() || t.isBF16()) { + return true; + } if (auto it = mlir::dyn_cast(t)) return (it.getWidth() == 8 || it.getWidth() == 16 || it.getWidth() == 32); return false; @@ -13690,7 +13700,6 @@ LogicalResult mlir::pto::TGemvAccOp::verify() { namespace mlir { namespace pto { - static LogicalResult parseShapeAndElem(AsmParser &parser, SmallVectorImpl &shape, Type &elementType, @@ -13890,7 +13899,9 @@ static void decomposeStridedLayout(AffineMap map, SmallVectorImpl &stri // 1. 初始化 strides.assign(map.getNumDims(), 0); - if (map.getNumResults() != 1) return; + if (map.getNumResults() != 1) { + return; + } // 2. 摊平表达式 SmallVector terms; @@ -13975,46 +13986,64 @@ static AffineMap buildStrictBitwiseAffineMap(MLIRContext *ctx, // Helper for parsing [64, 1] static ParseResult parseStrideList(AsmParser &parser, SmallVectorImpl &strides) { - if (parser.parseLSquare()) return failure(); + if (parser.parseLSquare()) { + return failure(); + } do { int64_t stride; - if (parser.parseInteger(stride)) return failure(); + if (parser.parseInteger(stride)) { + return failure(); + } strides.push_back(stride); } while (succeeded(parser.parseOptionalComma())); - if (parser.parseRSquare()) return failure(); + if (parser.parseRSquare()) { + return failure(); + } return success(); } // The custom attribute parser for: strided<[64, 1], offset: [?, ?]> [[maybe_unused]] static ParseResult parseStridedLayout(AsmParser &parser, Attribute &layout) { - if (parser.parseLess()) return failure(); + if (parser.parseLess()) { + return failure(); + } // 1. Parse Strides SmallVector strides; - if (parseStrideList(parser, strides)) return failure(); + if (parseStrideList(parser, strides)) { + return failure(); + } bool isMultiDim = false; unsigned numSymbols = 0; // 2. Parse Offset if (succeeded(parser.parseOptionalComma())) { - if (parser.parseKeyword("offset") || parser.parseColon()) return failure(); + if (parser.parseKeyword("offset") || parser.parseColon()) { + return failure(); + } // Check for multi-dim syntax: [?, ?] if (succeeded(parser.parseOptionalLSquare())) { isMultiDim = true; do { - if (parser.parseQuestion()) return failure(); + if (parser.parseQuestion()) { + return failure(); + } numSymbols++; } while (succeeded(parser.parseOptionalComma())); - if (parser.parseRSquare()) return failure(); + if (parser.parseRSquare()) { + return failure(); + } } else { // Fallback for old scalar syntax '?' if (parser.parseOptionalQuestion()) { /* handle single scalar */ } } } - if (parser.parseGreater()) return failure(); + if (parser.parseGreater()) { + return failure(); + } // 3. Validation if (isMultiDim && numSymbols != strides.size()) { @@ -14036,12 +14065,16 @@ static ParseResult parseStrideList(AsmParser &parser, SmallVectorImpl & // ============================================================================= [[maybe_unused]] static void printLayout(AsmPrinter &printer, Attribute layoutAttr) { - if (!layoutAttr) return; + if (!layoutAttr) { + return; + } auto mapAttr = llvm::dyn_cast(layoutAttr); if (!mapAttr) { printer << ", " << layoutAttr; return; } AffineMap map = mapAttr.getValue(); - if (map.isIdentity()) return; + if (map.isIdentity()) { + return; + } // 1. [核心修改] 反解 Strides SmallVector strides; @@ -14058,7 +14091,9 @@ static ParseResult parseStrideList(AsmParser &parser, SmallVectorImpl & printer << ", offset: ["; for (unsigned i = 0; i < numSyms; ++i) { printer << "?"; - if (i < numSyms - 1) printer << ", "; + if (i < numSyms - 1) { + printer << ", "; + } } printer << "]"; } @@ -14199,22 +14234,29 @@ LogicalResult SubViewOp::inferReturnTypes( MLIRContext *context, std::optional location, ValueRange operands, DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions, SmallVectorImpl &inferredReturnTypes) { - // 1. 获取 Source Type - if (operands.empty()) return failure(); + if (operands.empty()) { + return failure(); + } auto sourceType = llvm::dyn_cast(operands[0].getType()); - if (!sourceType) return failure(); + if (!sourceType) { + return failure(); + } // 2. 获取 subview 逻辑窗口（sizes） ArrayAttr sizeAttr; if (properties) { const auto *prop = properties.as(); - if (prop) sizeAttr = prop->sizes; + if (prop) { + sizeAttr = prop->sizes; + } } if (!sizeAttr && attributes) { sizeAttr = attributes.getAs("sizes"); } - if (!sizeAttr) return failure(); + if (!sizeAttr) { + return failure(); + } SmallVector subviewShape; for (auto attr : sizeAttr) { @@ -14279,7 +14321,9 @@ LogicalResult SubViewOp::inferReturnTypes( // 3. 继承 Config (若为空使用默认) auto cfg = sourceType.getConfigAttr(); - if (!cfg) cfg = TileBufConfigAttr::getDefault(context); + if (!cfg) { + cfg = TileBufConfigAttr::getDefault(context); + } // 4. 构建 Result Type auto canonicalValidShape = canonicalizeTileBufValidShape(validShape); @@ -14325,22 +14369,22 @@ static LogicalResult computeInnerShape(TileBufConfigAttr cfg, Type elemTy, bool &boxed, int32_t &bl, int32_t &sl) { auto readBLayoutI32 = [](Attribute attr, int32_t &out) -> bool { if (auto a = dyn_cast(attr)) { - out = (int32_t)a.getValue(); + out = static_cast(a.getValue()); return true; } if (auto a = dyn_cast(attr)) { - out = (int32_t)a.getInt(); + out = static_cast(a.getInt()); return true; } return false; }; auto readSLayoutI32 = [](Attribute attr, int32_t &out) -> bool { if (auto a = dyn_cast(attr)) { - out = (int32_t)a.getValue(); + out = static_cast(a.getValue()); return true; } if (auto a = dyn_cast(attr)) { - out = (int32_t)a.getInt(); + out = static_cast(a.getInt()); return true; } return false; @@ -14350,7 +14394,9 @@ static LogicalResult computeInnerShape(TileBufConfigAttr cfg, Type elemTy, int32_t fr = 512; (void)readBLayoutI32(cfg.getBLayout(), bl); (void)readSLayoutI32(cfg.getSLayout(), sl); - if (auto attr = dyn_cast(cfg.getSFractalSize())) fr = (int32_t)attr.getInt(); + if (auto attr = dyn_cast(cfg.getSFractalSize())) { + fr = static_cast(attr.getInt()); + } boxed = (sl != 0); if (!boxed) { @@ -14360,7 +14406,9 @@ static LogicalResult computeInnerShape(TileBufConfigAttr cfg, Type elemTy, } int64_t elemBytes = static_cast(getElemByteSize(elemTy)); - if (elemBytes <= 0) return failure(); + if (elemBytes <= 0) { + return failure(); + } if (fr == 1024) { innerRows = 16; @@ -14449,9 +14497,13 @@ mlir::LogicalResult mlir::pto::SubViewOp::verify() { if (dstTy.getMemorySpace() != srcTy.getMemorySpace()) return emitOpError("expects result address space to match source"); auto srcCfg = srcTy.getConfigAttr(); - if (!srcCfg) srcCfg = TileBufConfigAttr::getDefault(getContext()); + if (!srcCfg) { + srcCfg = TileBufConfigAttr::getDefault(getContext()); + } auto dstCfg = dstTy.getConfigAttr(); - if (!dstCfg) dstCfg = TileBufConfigAttr::getDefault(getContext()); + if (!dstCfg) { + dstCfg = TileBufConfigAttr::getDefault(getContext()); + } if (dstCfg != srcCfg) return emitOpError("expects result tile config to match source"); @@ -14487,7 +14539,9 @@ mlir::LogicalResult mlir::pto::SubViewOp::verify() { return emitOpError("expects result valid_shape[1] to match inferred/explicit valid_col"); auto cfg = srcTy.getConfigAttr(); - if (!cfg) cfg = TileBufConfigAttr::getDefault(getContext()); + if (!cfg) { + cfg = TileBufConfigAttr::getDefault(getContext()); + } int64_t innerRows = 1, innerCols = 1; bool boxed = false; @@ -18183,35 +18237,42 @@ static func::FuncOp getParentFunc(Operation *op) { static constexpr int64_t kSimtKeepResumeSlotLimit = 123; static Operation *getFirstNonConstantLikeOp(Block *block) { - if (!block) + if (!block) { return nullptr; + } for (Operation &op : *block) { - if (!op.hasTrait()) + if (!op.hasTrait()) { return &op; + } } return nullptr; } static bool isOpInRange(Operation *op, Operation *first, Operation *last) { for (Operation *cur = first; cur; cur = cur->getNextNode()) { - if (cur == op) + if (cur == op) { return true; - if (cur == last) + } + if (cur == last) { return false; + } } return false; } static std::optional getSimtKeepResumeRegisterCount(Type type) { if (auto intType = dyn_cast(type)) { - if (intType.getWidth() <= 32) + if (intType.getWidth() <= 32) { return 1; - if (intType.getWidth() == 64) + } + if (intType.getWidth() == 64) { return 2; + } return std::nullopt; } - if (type.isF16() || type.isBF16() || type.isF32()) + if (type.isF16() || type.isBF16() || type.isF32()) { return 1; + } return std::nullopt; } @@ -18232,22 +18293,26 @@ template static LogicalResult verifySimtKeepResumeSlotRange(OpT op) { std::optional registerCount = getSimtKeepResumeRegisterCount(getSimtKeepResumeValueType(op)); - if (!registerCount) + if (!registerCount) { return success(); + } int64_t slot = op.getSlot(); - if (slot < 0 || slot >= kSimtKeepResumeSlotLimit) + if (slot < 0 || slot >= kSimtKeepResumeSlotLimit) { return op.emitOpError() << "requires slot in range [0, " << (kSimtKeepResumeSlotLimit - 1) << "]"; + } if (*registerCount == 2) { - if ((slot % 2) != 0) + if ((slot % 2) != 0) { return op.emitOpError() << "requires an even slot for 64-bit keep/resume values"; - if (slot + 1 >= kSimtKeepResumeSlotLimit) + } + if (slot + 1 >= kSimtKeepResumeSlotLimit) { return op.emitOpError() << "requires slot in range [0, " << (kSimtKeepResumeSlotLimit - 2) << "] for 64-bit keep/resume values"; + } } return success(); } @@ -18257,15 +18322,18 @@ static bool overlapsEarlierSimtKeepResumeSlotUse(OpT op, SmallVectorImpl &used) { std::optional registerCount = getSimtKeepResumeRegisterCount(getSimtKeepResumeValueType(op)); - if (!registerCount) + if (!registerCount) { return false; + } int64_t slot = op.getSlot(); for (int64_t word = slot; word < slot + *registerCount; ++word) { - if (llvm::is_contained(used, word)) + if (llvm::is_contained(used, word)) { return true; + } } - for (int64_t word = slot; word < slot + *registerCount; ++word) + for (int64_t word = slot; word < slot + *registerCount; ++word) { used.push_back(word); + } return false; } @@ -18274,13 +18342,15 @@ static LogicalResult verifyUniqueResumeGroupSlots(ResumeOp current, SmallVector slots; for (Operation *cur = first; cur; cur = cur->getNextNode()) { auto resume = dyn_cast(cur); - if (!resume) + if (!resume) { break; + } if (overlapsEarlierSimtKeepResumeSlotUse(resume, slots) && - resume.getOperation() == current.getOperation()) + resume.getOperation() == current.getOperation()) { return current.emitOpError() << "duplicates an earlier slot " << resume.getSlot() << " in the SIMT resume prologue group"; + } } return success(); } @@ -18291,22 +18361,26 @@ static LogicalResult verifyUniqueKeepGroupSlots(KeepOp current, SmallVector slots; for (Operation *cur = first; cur; cur = cur->getNextNode()) { auto keep = dyn_cast(cur); - if (!keep) + if (!keep) { break; + } if (overlapsEarlierSimtKeepResumeSlotUse(keep, slots) && - keep.getOperation() == current.getOperation()) + keep.getOperation() == current.getOperation()) { return current.emitOpError() << "duplicates an earlier slot " << keep.getSlot() << " in the SIMT keep epilogue group"; - if (cur == last) + } + if (cur == last) { break; + } } return success(); } static bool isSupportedSimtKeepResumeType(Type type) { - if (auto intType = dyn_cast(type)) + if (auto intType = dyn_cast(type)) { return intType.getWidth() <= 64; + } return type.isF16() || type.isBF16() || type.isF32(); } diff --git a/lib/PTO/IR/PTOAttrs.cpp b/lib/PTO/IR/PTOAttrs.cpp index 2d306074c8..c001068f77 100644 --- a/lib/PTO/IR/PTOAttrs.cpp +++ b/lib/PTO/IR/PTOAttrs.cpp @@ -55,11 +55,21 @@ bool TileBufConfigAttr::isDefault() const { } static int32_t getLayoutInt(Attribute a, int32_t def) { - if (auto bl = mlir::dyn_cast(a)) return static_cast(bl.getValue()); - if (auto sl = mlir::dyn_cast(a)) return static_cast(sl.getValue()); - if (auto pv = mlir::dyn_cast(a)) return static_cast(pv.getValue()); - if (auto cm = mlir::dyn_cast(a)) return static_cast(cm.getValue()); - if (auto ia = mlir::dyn_cast(a)) return static_cast(ia.getInt()); + if (auto bl = mlir::dyn_cast(a)) { + return static_cast(bl.getValue()); + } + if (auto sl = mlir::dyn_cast(a)) { + return static_cast(sl.getValue()); + } + if (auto pv = mlir::dyn_cast(a)) { + return static_cast(pv.getValue()); + } + if (auto cm = mlir::dyn_cast(a)) { + return static_cast(cm.getValue()); + } + if (auto ia = mlir::dyn_cast(a)) { + return static_cast(ia.getInt()); + } return def; } @@ -83,7 +93,7 @@ LogicalResult TileBufConfigAttr::verify(function_ref emitE if (!sFractalSize || !sFractalSize.getType().isInteger(kI32BitWidth)) return emitError() << "s_fractal_size must be i32", failure(); - int32_t s = (int32_t)sFractalSize.getInt(); + int32_t s = static_cast(sFractalSize.getInt()); if (s != kFractalMxSize && s != kFractalABSize && s != kFractalCSize) return emitError() << "unsupported s_fractal_size: " << s << ", must be one of {" @@ -112,24 +122,39 @@ LogicalResult TileBufConfigAttr::verify(function_ref emitE // Helper: parse Attribute and convert to BLayoutAttr/SLayoutAttr/PadValueAttr static BLayoutAttr toBLayoutAttr(MLIRContext *ctx, Attribute a) { - if (auto bl = mlir::dyn_cast(a)) return bl; - if (auto ia = mlir::dyn_cast(a)) return BLayoutAttr::get(ctx, static_cast(ia.getInt())); + if (auto bl = mlir::dyn_cast(a)) { + return bl; + } + if (auto ia = mlir::dyn_cast(a)) { + return BLayoutAttr::get(ctx, static_cast(ia.getInt())); + } return {}; } static SLayoutAttr toSLayoutAttr(MLIRContext *ctx, Attribute a) { - if (auto sl = mlir::dyn_cast(a)) return sl; - if (auto ia = mlir::dyn_cast(a)) return SLayoutAttr::get(ctx, static_cast(ia.getInt())); + if (auto sl = mlir::dyn_cast(a)) { + return sl; + } + if (auto ia = mlir::dyn_cast(a)) { + return SLayoutAttr::get(ctx, static_cast(ia.getInt())); + } return {}; } static PadValueAttr toPadValueAttr(MLIRContext *ctx, Attribute a) { - if (auto pv = mlir::dyn_cast(a)) return pv; - if (auto ia = mlir::dyn_cast(a)) return PadValueAttr::get(ctx, static_cast(ia.getInt())); + if (auto pv = mlir::dyn_cast(a)) { + return pv; + } + if (auto ia = mlir::dyn_cast(a)) { + return PadValueAttr::get(ctx, static_cast(ia.getInt())); + } return {}; } static CompactModeAttr toCompactModeAttr(MLIRContext *ctx, Attribute a) { - if (auto cm = mlir::dyn_cast(a)) return cm; - if (auto ia = mlir::dyn_cast(a)) + if (auto cm = mlir::dyn_cast(a)) { + return cm; + } + if (auto ia = mlir::dyn_cast(a)) { return CompactModeAttr::get(ctx, static_cast(ia.getInt())); + } return {}; } @@ -142,50 +167,78 @@ Attribute TileBufConfigAttr::parse(AsmParser &p, Type) { PadValueAttr pv = def.getPad(); CompactModeAttr compact = def.getCompactMode(); - if (p.parseLess()) return {}; + if (p.parseLess()) { + return {}; + } - if (succeeded(p.parseOptionalGreater())) + if (succeeded(p.parseOptionalGreater())) { return TileBufConfigAttr::get(ctx, bl, sl, sz, pv, compact); + } bool parsedGreater = false; while (!parsedGreater) { StringRef key; - if (p.parseKeyword(&key)) return {}; - if (p.parseEqual()) return {}; + if (p.parseKeyword(&key)) { + return {}; + } + if (p.parseEqual()) { + return {}; + } if (key == "blayout") { Attribute a; - if (p.parseAttribute(a)) return {}; + if (p.parseAttribute(a)) { + return {}; + } bl = toBLayoutAttr(ctx, a); - if (!bl) return {}; + if (!bl) { + return {}; + } } else if (key == "slayout") { Attribute a; - if (p.parseAttribute(a)) return {}; + if (p.parseAttribute(a)) { + return {}; + } sl = toSLayoutAttr(ctx, a); - if (!sl) return {}; + if (!sl) { + return {}; + } } else if (key == "s_fractal_size") { int32_t v = 0; - if (p.parseInteger(v)) return {}; + if (p.parseInteger(v)) { + return {}; + } sz = IntegerAttr::get(IntegerType::get(ctx, kI32BitWidth), v); } else if (key == "pad") { Attribute a; - if (p.parseAttribute(a)) return {}; + if (p.parseAttribute(a)) { + return {}; + } pv = toPadValueAttr(ctx, a); - if (!pv) return {}; + if (!pv) { + return {}; + } } else if (key == "compact") { Attribute a; - if (p.parseAttribute(a)) return {}; + if (p.parseAttribute(a)) { + return {}; + } compact = toCompactModeAttr(ctx, a); - if (!compact) return {}; + if (!compact) { + return {}; + } } else { p.emitError(p.getCurrentLocation(), "unknown key in tile_buf_config: ") << key; return {}; } parsedGreater = succeeded(p.parseOptionalGreater()); - if (parsedGreater) + if (parsedGreater) { break; - if (p.parseComma()) return {}; + } + if (p.parseComma()) { + return {}; + } } return TileBufConfigAttr::get(ctx, bl, sl, sz, pv, compact); @@ -195,7 +248,7 @@ void TileBufConfigAttr::print(AsmPrinter &p) const { p << "<"; p << "blayout=" << getBLayout(); p << ", slayout=" << getSLayout(); - p << ", s_fractal_size=" << (int32_t)getSFractalSize().getInt(); + p << ", s_fractal_size=" << static_cast(getSFractalSize().getInt()); p << ", pad=" << getPad(); p << ", compact=" << getCompactMode(); p << ">"; diff --git a/lib/PTO/IR/PTOTypeDefs.cpp b/lib/PTO/IR/PTOTypeDefs.cpp index c7594a7da7..b22958e2bc 100644 --- a/lib/PTO/IR/PTOTypeDefs.cpp +++ b/lib/PTO/IR/PTOTypeDefs.cpp @@ -33,8 +33,9 @@ using TileBufValidShapeVector = void mlir::pto::setPTOParserTargetArch(MLIRContext *context, PTOParserTargetArch arch) { - if (!context) + if (!context) { return; + } std::lock_guard lock(parserTargetArchMutex); if (arch == PTOParserTargetArch::Unspecified) { @@ -45,13 +46,15 @@ void mlir::pto::setPTOParserTargetArch(MLIRContext *context, } PTOParserTargetArch mlir::pto::getPTOParserTargetArch(MLIRContext *context) { - if (!context) + if (!context) { return PTOParserTargetArch::Unspecified; + } std::lock_guard lock(parserTargetArchMutex); auto it = parserTargetArchByContext.find(context); - if (it == parserTargetArchByContext.end()) + if (it == parserTargetArchByContext.end()) { return PTOParserTargetArch::Unspecified; + } return it->second; } @@ -69,8 +72,9 @@ static TileBufValidShapeVector canonicalizeTileBufValidShape(ArrayRef validShape) { TileBufValidShapeVector canonical; canonical.reserve(validShape.size()); - for (int64_t dim : validShape) + for (int64_t dim : validShape) { canonical.push_back(dim < 0 ? ShapedType::kDynamic : dim); + } return canonical; } @@ -174,8 +178,9 @@ static std::optional resolveTileBufMemorySpace(StringRef locStr) { static BLayout resolveTileBufBLayout(MLIRContext *context, AddressSpace memorySpace, BLayout parsedLayout) { - if (memorySpace != AddressSpace::LEFT) + if (memorySpace != AddressSpace::LEFT) { return parsedLayout; + } switch (getPTOParserTargetArch(context)) { case PTOParserTargetArch::A3: @@ -192,12 +197,16 @@ TileBufConfigAttr TileBufType::getConfigAttr() const { // 情况 A：getConfig() 已经是 TileBufConfigAttr if constexpr (std::is_same_v) { auto cfg = getConfig(); - if (!cfg) cfg = TileBufConfigAttr::getDefault(getContext()); + if (!cfg) { + cfg = TileBufConfigAttr::getDefault(getContext()); + } return cfg; } else { // 情况 B：getConfig() 是 Attribute auto cfg = llvm::dyn_cast_or_null(getConfig()); - if (!cfg) cfg = TileBufConfigAttr::getDefault(getContext()); + if (!cfg) { + cfg = TileBufConfigAttr::getDefault(getContext()); + } return cfg; } } @@ -214,7 +223,7 @@ mlir::Attribute TileBufType::getCompactModeAttr() const { // ✅ numeric getters（可选） int32_t TileBufType::getSFractalSizeI32() const { - return (int32_t)getConfigAttr().getSFractalSize().getInt(); + return static_cast(getConfigAttr().getSFractalSize().getInt()); } int32_t TileBufType::getBLayoutValueI32() const { @@ -637,8 +646,9 @@ void mlir::pto::TileBufType::print(mlir::AsmPrinter &printer) const { int64_t cols = shape.size() > 1 ? shape[1] : ShapedType::kDynamic; auto cfg = getConfigAttr(); - if (!cfg) + if (!cfg) { cfg = mlir::pto::TileBufConfigAttr::getDefault(getContext()); + } auto defaultCfg = TileBufConfigAttr::getDefault(getContext()); llvm::StringRef locStr = stringifyLocFromMemorySpace(getMemorySpace()); @@ -701,18 +711,6 @@ void mlir::pto::TileBufType::print(mlir::AsmPrinter &printer) const { } // ---- MultiTileBufType custom asm ---- -// -// Syntax: -// -// Verbose form: -// !pto.multi_tile_buf, count=N> -// -// Compact (sugar) form: -// !pto.multi_tile_buf -// -// In the compact form the per-slot tile_buf is built from the same compact -// syntax as `!pto.tile_buf`, followed by a mandatory `count=N`. - LogicalResult MultiTileBufType::verify( function_ref emitError, mlir::pto::TileBufType slotType, uint32_t count) { diff --git a/lib/PTO/IR/PTOTypeUtils.cpp b/lib/PTO/IR/PTOTypeUtils.cpp index db8aade8cd..dc65a30834 100644 --- a/lib/PTO/IR/PTOTypeUtils.cpp +++ b/lib/PTO/IR/PTOTypeUtils.cpp @@ -42,8 +42,9 @@ bool mlir::pto::isPTOFloat4PackedType(Type t) { bool mlir::pto::isPTOPackedLdgStgVectorType(Type t) { // !pto.hif8x2 is a 2-byte packed hif8 value type (not a VectorType). - if (isPTOHiFloat8x2Type(t)) + if (isPTOHiFloat8x2Type(t)) { return true; + } auto vecType = dyn_cast(t); if (!vecType || vecType.isScalable() || vecType.getRank() != 1) return false; @@ -63,8 +64,9 @@ bool mlir::pto::isPTOPackedLdgStgVectorType(Type t) { validElem = lanes == 2 && (w == 8 || w == 16 || w == 32); } } - if (!validElem) + if (!validElem) { return false; + } unsigned totalBits = vecType.getDimSize(0) * getPTOStorageElemBitWidth(elemType); return totalBits == 16 || totalBits == 32 || totalBits == 64; @@ -84,8 +86,9 @@ bool mlir::pto::isPTOLowPrecisionType(Type t) { } unsigned mlir::pto::getPTOStorageElemBitWidth(Type t) { - if (isPTOHiFloat8x2Type(t)) + if (isPTOHiFloat8x2Type(t)) { return 16; + } if (isPTOLowPrecisionType(t)) return kBitsPerByte; if (auto floatTy = dyn_cast(t)) diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp index 6f5143959e..bc9697c2e1 100644 --- a/lib/PTO/IR/VMI.cpp +++ b/lib/PTO/IR/VMI.cpp @@ -26,14 +26,14 @@ using namespace mlir; using namespace mlir::pto; namespace { - static std::string formatVMIVRegType(int64_t elementCount, Type elementType, Attribute layout) { std::string result; llvm::raw_string_ostream os(result); os << "!pto.vmi.vreg<" << elementCount << "x" << elementType; - if (layout) + if (layout) { os << ", " << layout; + } os << ">"; return result; } @@ -43,8 +43,9 @@ static std::string formatVMIMaskType(int64_t elementCount, std::string result; llvm::raw_string_ostream os(result); os << "!pto.vmi.mask<" << elementCount << "x" << granularity; - if (layout) + if (layout) { os << ", " << layout; + } os << ">"; return result; } @@ -77,16 +78,18 @@ static bool isVMIPredicateMaskableElementType(Type type) { static bool isVMIAnyI8I16I32Type(Type type) { auto integerType = dyn_cast(type); - if (!integerType) + if (!integerType) { return false; + } return integerType.getWidth() == 8 || integerType.getWidth() == 16 || integerType.getWidth() == 32; } static bool isVMISignedOrSignlessI8I16I32Type(Type type) { auto integerType = dyn_cast(type); - if (!integerType || integerType.isUnsigned()) + if (!integerType || integerType.isUnsigned()) { return false; + } return integerType.getWidth() == 8 || integerType.getWidth() == 16 || integerType.getWidth() == 32; } @@ -142,8 +145,9 @@ static bool isVMIIotaElementType(Type type) { static bool isCompatibleScalarForSemanticType(Type semanticType, Type scalarType) { - if (semanticType == scalarType) + if (semanticType == scalarType) { return true; + } auto semanticInt = dyn_cast(semanticType); auto scalarInt = dyn_cast(scalarType); @@ -151,24 +155,29 @@ static bool isCompatibleScalarForSemanticType(Type semanticType, semanticInt.getWidth() != scalarInt.getWidth()) return false; - if (semanticInt.isSigned()) + if (semanticInt.isSigned()) { return scalarInt.isSigned() || scalarInt.isSignless(); - if (semanticInt.isUnsigned()) + } + if (semanticInt.isUnsigned()) { return scalarInt.isUnsigned() || scalarInt.isSignless(); + } return scalarInt.isSignless(); } static unsigned getVMIElementBitWidth(Type type) { - if (isa(type)) + if (isa(type)) { return 64; + } return pto::getPTOStorageElemBitWidth(type); } static std::optional getVMIIntegerOrFloatBitWidth(Type type) { - if (auto intType = dyn_cast(type)) + if (auto intType = dyn_cast(type)) { return intType.getWidth(); - if (auto floatType = dyn_cast(type)) + } + if (auto floatType = dyn_cast(type)) { return floatType.getWidth(); + } return std::nullopt; } @@ -178,11 +187,13 @@ static int64_t divideCeilNonNegative(int64_t value, int64_t divisor) { static LogicalResult parseOptionalVMILayout(AsmParser &parser, Attribute &layout) { - if (failed(parser.parseOptionalComma())) + if (failed(parser.parseOptionalComma())) { return success(); + } - if (failed(parser.parseAttribute(layout))) + if (failed(parser.parseAttribute(layout))) { return failure(); + } if (!mlir::isa(layout)) return parser.emitError(parser.getCurrentLocation(), "expected #pto.vmi.layout attribute"); @@ -190,32 +201,39 @@ static LogicalResult parseOptionalVMILayout(AsmParser &parser, } static FailureOr getVMIElementCount(Type type) { - if (auto vregType = dyn_cast(type)) + if (auto vregType = dyn_cast(type)) { return vregType.getElementCount(); - if (auto maskType = dyn_cast(type)) + } + if (auto maskType = dyn_cast(type)) { return maskType.getElementCount(); + } return failure(); } static FailureOr getAssignedVMILayout(Type type) { Attribute layout; - if (auto vregType = dyn_cast(type)) + if (auto vregType = dyn_cast(type)) { layout = vregType.getLayout(); - else if (auto maskType = dyn_cast(type)) + } + else if (auto maskType = dyn_cast(type)) { layout = maskType.getLayout(); - else + } + else { return failure(); + } auto layoutAttr = dyn_cast_or_null(layout); - if (!layoutAttr) + if (!layoutAttr) { return failure(); + } return layoutAttr; } static FailureOr getLayoutFactor(Type type) { FailureOr layout = getAssignedVMILayout(type); - if (failed(layout)) + if (failed(layout)) { return failure(); + } return (*layout).isDenseSplit() ? (*layout).getFactor() : 1; } @@ -224,12 +242,15 @@ static FailureOr getLayoutBlockElems(Type type) { } static int64_t getMaskGranularityBitWidth(StringRef granularity) { - if (granularity == "b8") + if (granularity == "b8") { return 8; - if (granularity == "b16") + } + if (granularity == "b16") { return 16; - if (granularity == "b32") + } + if (granularity == "b32") { return 32; + } return 0; } @@ -248,16 +269,18 @@ static StringRef getMaskGranularityForBitWidth(int64_t bits) { static FailureOr getVMIMaskPhysicalGranularity(VMIMaskType type) { int64_t bits = getMaskGranularityBitWidth(type.getGranularity()); - if (bits == 0) + if (bits == 0) { return failure(); + } VMILayoutAttr layout = type.getLayoutAttr(); int64_t laneStride = layout && layout.hasLaneStride() ? layout.getLaneStride() : 1; StringRef physicalGranularity = getMaskGranularityForBitWidth(bits * laneStride); - if (physicalGranularity.empty()) + if (physicalGranularity.empty()) { return failure(); + } return physicalGranularity; } @@ -268,8 +291,9 @@ static FailureOr getPhysicalLanesPerPart(Type type) { if (auto maskType = dyn_cast(type)) { FailureOr physicalGranularity = getVMIMaskPhysicalGranularity(maskType); - if (failed(physicalGranularity)) + if (failed(physicalGranularity)) { return failure(); + } return getMaskLanesPerPart(*physicalGranularity); } return failure(); @@ -277,10 +301,12 @@ static FailureOr getPhysicalLanesPerPart(Type type) { static FailureOr getDenseLaneStride(Type type) { FailureOr layout = getAssignedVMILayout(type); - if (failed(layout)) + if (failed(layout)) { return failure(); - if (isa(type)) + } + if (isa(type)) { return 1; + } return (*layout).isDense() ? (*layout).getLaneStride() : 1; } @@ -295,8 +321,9 @@ static bool isLayoutAssigned(VMIMaskType type) { static LogicalResult verifyAllSameVRegShapeAndLayout(Operation *op, ArrayRef types, bool requireSameElement) { - if (types.empty()) + if (types.empty()) { return success(); + } VMIVRegType first = types.front(); bool anyLayout = llvm::any_of( @@ -321,8 +348,9 @@ verifyAllSameVRegShapeAndLayout(Operation *op, ArrayRef types, static LogicalResult verifyAllSameVRegShapeAndLayoutPresence( Operation *op, ArrayRef types, bool requireSameElement) { - if (types.empty()) + if (types.empty()) { return success(); + } VMIVRegType first = types.front(); bool anyLayout = llvm::any_of( @@ -351,8 +379,9 @@ static LogicalResult verifyElementwiseVRegOp(Operation *op, VMIVRegType lhs, static LogicalResult verifyFloatUnaryVRegOp(Operation *op, VMIVRegType source, VMIVRegType result) { - if (!isVMIFloatLikeType(source.getElementType())) + if (!isVMIFloatLikeType(source.getElementType())) { return op->emitOpError("requires floating-point-like VMI element type"); + } return verifyAllSameVRegShapeAndLayout(op, {source, result}, /*requireSameElement=*/true); } @@ -360,8 +389,9 @@ static LogicalResult verifyFloatUnaryVRegOp(Operation *op, VMIVRegType source, static LogicalResult verifyFloatTernaryVRegOp(Operation *op, VMIVRegType lhs, VMIVRegType rhs, VMIVRegType acc, VMIVRegType result) { - if (!isVMIFloatLikeType(lhs.getElementType())) + if (!isVMIFloatLikeType(lhs.getElementType())) { return op->emitOpError("requires floating-point-like VMI element type"); + } return verifyAllSameVRegShapeAndLayout(op, {lhs, rhs, acc, result}, /*requireSameElement=*/true); } @@ -369,8 +399,9 @@ static LogicalResult verifyFloatTernaryVRegOp(Operation *op, VMIVRegType lhs, static LogicalResult verifyAllSameMaskShapeLayoutAndGranularity(Operation *op, ArrayRef types) { - if (types.empty()) + if (types.empty()) { return success(); + } VMIMaskType first = types.front(); bool anyLayout = llvm::any_of( @@ -404,12 +435,14 @@ static LogicalResult verifyMaskMatchesData(Operation *op, VMIMaskType maskType, if (!isLayoutAssigned(maskType) || !isLayoutAssigned(dataType)) return op->emitOpError("requires either both mask and data to carry " "layout or neither to carry layout"); - if (maskType.getLayout() != dataType.getLayout()) + if (maskType.getLayout() != dataType.getLayout()) { return op->emitOpError("requires mask layout to match data layout"); + } } - if (maskType.isPred()) + if (maskType.isPred()) { return success(); + } unsigned elementBitWidth = getVMIElementBitWidth(dataType.getElementType()); int64_t maskBitWidth = getMaskGranularityBitWidth(maskType.getGranularity()); @@ -423,33 +456,40 @@ static LogicalResult verifyMaskMatchesData(Operation *op, VMIMaskType maskType, static Type getMemoryElementType(Type type) { - if (auto ptrType = dyn_cast(type)) + if (auto ptrType = dyn_cast(type)) { return ptrType.getElementType(); - if (auto memrefType = dyn_cast(type)) + } + if (auto memrefType = dyn_cast(type)) { return memrefType.getElementType(); + } return {}; } static bool isUBBackedMemoryType(Type type) { - if (auto ptrType = dyn_cast(type)) + if (auto ptrType = dyn_cast(type)) { return ptrType.getMemorySpace().getAddressSpace() == AddressSpace::VEC; + } auto memrefType = dyn_cast(type); - if (!memrefType) + if (!memrefType) { return false; + } Attribute memorySpace = memrefType.getMemorySpace(); - if (auto addressSpace = dyn_cast_or_null(memorySpace)) + if (auto addressSpace = dyn_cast_or_null(memorySpace)) { return addressSpace.getAddressSpace() == AddressSpace::VEC; - if (auto integerSpace = dyn_cast_or_null(memorySpace)) + } + if (auto integerSpace = dyn_cast_or_null(memorySpace)) { return integerSpace.getInt() == static_cast(AddressSpace::VEC); + } return false; } static LogicalResult verifyUBBackedMemory(Operation *op, Type memoryType, StringRef role) { - if (isUBBackedMemoryType(memoryType)) + if (isUBBackedMemoryType(memoryType)) { return success(); + } return op->emitOpError() << "requires memory " << role << " to be UB-backed"; } @@ -458,8 +498,9 @@ static LogicalResult verifyMemoryElementMatches(Operation *op, Type memoryType, VMIVRegType dataType, StringRef role) { Type memoryElementType = getMemoryElementType(memoryType); - if (!memoryElementType) + if (!memoryElementType) { return success(); + } if (memoryElementType != dataType.getElementType()) return op->emitOpError() << "requires memory " << role << " element type to match VMI data element type"; @@ -479,8 +520,9 @@ static LogicalResult verifyContiguousIfLayoutAssigned(Operation *op, static bool isPackedByteGroupStore(Type memoryType, VMIVRegType dataType) { Type memoryElementType = getMemoryElementType(memoryType); - if (!memoryElementType) + if (!memoryElementType) { return false; + } auto memoryIntegerType = dyn_cast(memoryElementType); auto dataIntegerType = dyn_cast(dataType.getElementType()); return memoryIntegerType && dataIntegerType && @@ -489,8 +531,9 @@ static bool isPackedByteGroupStore(Type memoryType, VMIVRegType dataType) { static LogicalResult verifyNumGroups(Operation *op, VMIVRegType type, int64_t numGroups) { - if (numGroups <= 0) + if (numGroups <= 0) { return op->emitOpError("requires num_groups to be positive"); + } if (type.getElementCount() % numGroups != 0) return op->emitOpError() << "requires num_groups to evenly divide VMI logical lane count " @@ -517,8 +560,9 @@ static LogicalResult verifyPhysicalParts(Operation *op, Type vmiType, "requires data element type with known physical lane count"); for (Type physicalType : physicalTypes) { auto partType = dyn_cast(physicalType); - if (!partType) + if (!partType) { return op->emitOpError("requires physical data parts to be !pto.vreg"); + } if (partType.getElementCount() != *lanesPerPart || partType.getElementType() != physicalElementType) return op->emitOpError( @@ -528,8 +572,9 @@ static LogicalResult verifyPhysicalParts(Operation *op, Type vmiType, } auto maskType = dyn_cast(vmiType); - if (!maskType) + if (!maskType) { return op->emitOpError("requires VMI data or mask type"); + } if (maskType.isPred()) return op->emitOpError( "requires layout-assigned mask with concrete granularity"); @@ -541,8 +586,9 @@ static LogicalResult verifyPhysicalParts(Operation *op, Type vmiType, for (Type physicalType : physicalTypes) { auto partType = dyn_cast(physicalType); - if (!partType) + if (!partType) { return op->emitOpError("requires physical mask parts to be !pto.mask"); + } if (partType.getGranularity() != *physicalGranularity) return op->emitOpError( "requires physical mask part granularity to match VMI mask carrier"); @@ -575,8 +621,9 @@ mapDensePartIndexToLogicalLane(int64_t elementCount, int64_t factor, int64_t inBlockLane = indexInPart % blockElems; int64_t logicalBlock = partBlock * factor + part; int64_t logicalLane = logicalBlock * blockElems + inBlockLane; - if (logicalLane >= elementCount) + if (logicalLane >= elementCount) { return std::nullopt; + } return logicalLane; } @@ -587,8 +634,9 @@ static int64_t getDenseLogicalLanesInPart(int64_t elementCount, int64_t factor, int64_t lanePart = 0; std::optional index = mapDenseLogicalLaneToPartIndex( elementCount, factor, blockElems, lane, lanePart); - if (index && lanePart == part) + if (index && lanePart == part) { maxIndex = std::max(maxIndex, *index); + } } return maxIndex + 1; } @@ -600,15 +648,16 @@ static int64_t getDenseLogicalLanesInPart(int64_t elementCount, int64_t factor, // --------------------------------------------------------------------------- namespace mlir::pto { - std::optional lookupVMIFpToSiContract(Type srcElem, Type dstElem) { // Must be float → signed/signless integer. - if (!isVMIFloatLikeType(srcElem)) + if (!isVMIFloatLikeType(srcElem)) { return std::nullopt; + } auto dstInt = dyn_cast(dstElem); - if (!dstInt || dstInt.isUnsigned()) + if (!dstInt || dstInt.isUnsigned()) { return std::nullopt; + } bool srcF32 = srcElem.isF32(); bool srcF16 = srcElem.isF16(); @@ -644,11 +693,13 @@ lookupVMIFpToSiContract(Type srcElem, Type dstElem) { std::optional lookupVMIFpToUIContract(Type srcElem, Type dstElem) { // Must be float → unsigned integer. - if (!isVMIFloatLikeType(srcElem)) + if (!isVMIFloatLikeType(srcElem)) { return std::nullopt; + } auto dstInt = dyn_cast(dstElem); - if (!dstInt || !dstInt.isUnsigned()) + if (!dstInt || !dstInt.isUnsigned()) { return std::nullopt; + } bool srcF16 = srcElem.isF16(); unsigned dstBits = dstInt.getWidth(); @@ -667,12 +718,14 @@ lookupVMIFpToUIContract(Type srcElem, Type dstElem) { std::optional lookupVMIFpToFpContract(Type srcElem, Type dstElem) { - if (!isVMIFloatLikeType(srcElem) || !isVMIFloatLikeType(dstElem)) + if (!isVMIFloatLikeType(srcElem) || !isVMIFloatLikeType(dstElem)) { return std::nullopt; + } unsigned srcBits = pto::getPTOStorageElemBitWidth(srcElem); unsigned dstBits = pto::getPTOStorageElemBitWidth(dstElem); - if (srcBits != dstBits) + if (srcBits != dstBits) { return std::nullopt; + } // bf16 -> f16: same-width, rnd, sat, no part. if (srcElem.isBF16() && dstElem.isF16()) return VMIFpToFpContract{/*requiresRnd=*/true, /*requiresSat=*/true, @@ -714,8 +767,9 @@ Attribute VMILayoutAttr::parse(AsmParser &parser, Type) { int64_t slots = 0; int64_t laneStride = 1; - if (failed(parser.parseLess()) || failed(parser.parseKeyword(&kind))) + if (failed(parser.parseLess()) || failed(parser.parseKeyword(&kind))) { return {}; + } if (kind == "contiguous") { factor = 1; @@ -729,15 +783,18 @@ Attribute VMILayoutAttr::parse(AsmParser &parser, Type) { } } } else if (kind == "deinterleaved") { - if (failed(parser.parseEqual()) || failed(parser.parseInteger(factor))) + if (failed(parser.parseEqual()) || failed(parser.parseInteger(factor))) { return {}; + } while (succeeded(parser.parseOptionalComma())) { StringRef field; - if (failed(parser.parseKeyword(&field)) || failed(parser.parseEqual())) + if (failed(parser.parseKeyword(&field)) || failed(parser.parseEqual())) { return {}; + } if (field == "lane_stride") { - if (failed(parser.parseInteger(laneStride))) + if (failed(parser.parseInteger(laneStride))) { return {}; + } } else { parser.emitError(parser.getCurrentLocation(), "expected 'lane_stride = '"); @@ -745,21 +802,26 @@ Attribute VMILayoutAttr::parse(AsmParser &parser, Type) { } } } else if (kind == "block_deinterleaved") { - if (failed(parser.parseEqual()) || failed(parser.parseInteger(factor))) + if (failed(parser.parseEqual()) || failed(parser.parseInteger(factor))) { return {}; + } } else if (kind == "num_groups") { - if (failed(parser.parseEqual()) || failed(parser.parseInteger(factor))) + if (failed(parser.parseEqual()) || failed(parser.parseInteger(factor))) { return {}; + } while (succeeded(parser.parseOptionalComma())) { StringRef field; - if (failed(parser.parseKeyword(&field)) || failed(parser.parseEqual())) + if (failed(parser.parseKeyword(&field)) || failed(parser.parseEqual())) { return {}; + } if (field == "slots") { - if (failed(parser.parseInteger(slots))) + if (failed(parser.parseInteger(slots))) { return {}; + } } else if (field == "lane_stride") { - if (failed(parser.parseInteger(laneStride))) + if (failed(parser.parseInteger(laneStride))) { return {}; + } } else { parser.emitError(parser.getCurrentLocation(), "expected 'slots = ' or " @@ -775,8 +837,9 @@ Attribute VMILayoutAttr::parse(AsmParser &parser, Type) { return {}; } - if (failed(parser.parseGreater())) + if (failed(parser.parseGreater())) { return {}; + } return parser.getChecked(loc, parser.getContext(), kind, factor, blockElems, slots, @@ -786,20 +849,24 @@ Attribute VMILayoutAttr::parse(AsmParser &parser, Type) { void VMILayoutAttr::print(AsmPrinter &printer) const { printer << "<" << getKind(); if (isContiguous()) { - if (getLaneStride() != 1) + if (getLaneStride() != 1) { printer << ", lane_stride = " << getLaneStride(); + } } else if (isDeinterleaved()) { printer << " = " << getFactor(); - if (getLaneStride() != 1) + if (getLaneStride() != 1) { printer << ", lane_stride = " << getLaneStride(); + } } else if (isBlockDeinterleaved()) { printer << " = " << getFactor(); } else if (isGroupSlots()) { printer << " = " << getFactor(); - if (getSlots() != 0) + if (getSlots() != 0) { printer << ", slots = " << getSlots(); - if (getLaneStride() != 1) + } + if (getLaneStride() != 1) { printer << ", lane_stride = " << getLaneStride(); + } } printer << ">"; } @@ -885,8 +952,9 @@ Type VMIVRegType::parse(AsmParser &parser) { void VMIVRegType::print(AsmPrinter &printer) const { printer << "<" << getElementCount() << "x"; printer.printType(getElementType()); - if (getLayout()) + if (getLayout()) { printer << ", " << getLayout(); + } printer << ">"; } @@ -960,8 +1028,9 @@ Type VMIMaskType::parse(AsmParser &parser) { void VMIMaskType::print(AsmPrinter &printer) const { printer << "<" << getElementCount() << "x" << getGranularity(); - if (getLayout()) + if (getLayout()) { printer << ", " << getLayout(); + } printer << ">"; } @@ -1001,8 +1070,9 @@ LogicalResult VMIMaskType::verify(function_ref emitError, LogicalResult VMIConstantOp::verify() { auto resultType = cast(getResult().getType()); auto denseAttr = dyn_cast(getValue()); - if (!denseAttr) + if (!denseAttr) { return emitOpError("requires dense elements constant attribute"); + } if (denseAttr.getElementType() != resultType.getElementType()) return emitOpError( "requires dense constant element type to match result element type"); @@ -1015,11 +1085,13 @@ LogicalResult VMIConstantOp::verify() { LogicalResult VMIBroadcastOp::verify() { auto resultType = cast(getResult().getType()); Type valueType = getValue().getType(); - if (valueType == resultType.getElementType()) + if (valueType == resultType.getElementType()) { return success(); + } if (auto vregType = dyn_cast(valueType)) { - if (vregType.getElementCount() != 1) + if (vregType.getElementCount() != 1) { return emitOpError("requires VMI vector input to have one logical lane"); + } if (vregType.getElementType() != resultType.getElementType()) return emitOpError("requires VMI vector input element type to match " "result element type"); @@ -1035,12 +1107,14 @@ LogicalResult VMIIotaOp::verify() { if (!isVMIIotaElementType(elementType)) return emitOpError("requires result element type to be integer 8/16/32 " "or f16/f32"); - if (!isCompatibleScalarForSemanticType(elementType, getBase().getType())) + if (!isCompatibleScalarForSemanticType(elementType, getBase().getType())) { return emitOpError("requires base type to match result element type"); + } if (std::optional order = getOrder()) { - if (*order != "ASC" && *order != "DESC") + if (*order != "ASC" && *order != "DESC") { return emitOpError("requires order to be ASC or DESC"); + } } return success(); } @@ -1051,16 +1125,19 @@ LogicalResult VMIGroupIotaOp::verify() { if (!isVMIIotaElementType(elementType)) return emitOpError("requires result element type to be integer 8/16/32 " "or f16/f32"); - if (!isCompatibleScalarForSemanticType(elementType, getBase().getType())) + if (!isCompatibleScalarForSemanticType(elementType, getBase().getType())) { return emitOpError("requires base type to match result element type"); + } if (std::optional order = getOrder()) { - if (*order != "ASC" && *order != "DESC") + if (*order != "ASC" && *order != "DESC") { return emitOpError("requires order to be ASC or DESC"); + } } int64_t numGroups = getGroupAttr().getInt(); - if (numGroups <= 1) + if (numGroups <= 1) { return emitOpError("requires group greater than one"); + } if (resultType.getElementCount() % numGroups != 0) return emitOpError("requires group to evenly divide result logical lane " "count"); @@ -1085,10 +1162,12 @@ LogicalResult VMICreateGroupMaskOp::verify() { auto resultType = cast(getResult().getType()); int64_t numGroups = getNumGroupsAttr().getInt(); int64_t groupSize = getGroupSizeAttr().getInt(); - if (numGroups <= 0) + if (numGroups <= 0) { return emitOpError("requires positive num_groups"); - if (groupSize <= 0) + } + if (groupSize <= 0) { return emitOpError("requires positive group_size"); + } if (resultType.getElementCount() != numGroups * groupSize) return emitOpError("requires result lane count to equal num_groups * " "group_size"); @@ -1098,10 +1177,12 @@ LogicalResult VMICreateGroupMaskOp::verify() { LogicalResult VMIConstantMaskOp::verify() { auto resultType = cast(getResult().getType()); auto denseAttr = dyn_cast(getValue()); - if (!denseAttr) + if (!denseAttr) { return emitOpError("requires dense elements mask constant attribute"); - if (!denseAttr.getElementType().isInteger(1)) + } + if (!denseAttr.getElementType().isInteger(1)) { return emitOpError("requires dense mask constant element type to be i1"); + } if (denseAttr.getNumElements() != resultType.getElementCount()) return emitOpError("requires dense mask constant element count to match " "result logical lane count"); @@ -1143,10 +1224,12 @@ LogicalResult VMIAddFOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(lhsType.getElementType())) + if (!isVMIFloatLikeType(lhsType.getElementType())) { return emitOpError("requires floating-point-like VMI element type"); - if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + } + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) { return emitOpError("requires f16, bf16, or f32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1154,10 +1237,12 @@ LogicalResult VMIAddIOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(lhsType.getElementType())) + if (!isVMIIntegerLikeType(lhsType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + } + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) { return emitOpError("requires i8, i16, or i32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1165,10 +1250,12 @@ LogicalResult VMISubFOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(lhsType.getElementType())) + if (!isVMIFloatLikeType(lhsType.getElementType())) { return emitOpError("requires floating-point-like VMI element type"); - if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + } + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) { return emitOpError("requires f16, bf16, or f32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1176,10 +1263,12 @@ LogicalResult VMISubIOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(lhsType.getElementType())) + if (!isVMIIntegerLikeType(lhsType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + } + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) { return emitOpError("requires i8, i16, or i32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1187,10 +1276,12 @@ LogicalResult VMIMulFOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(lhsType.getElementType())) + if (!isVMIFloatLikeType(lhsType.getElementType())) { return emitOpError("requires floating-point-like VMI element type"); - if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + } + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) { return emitOpError("requires f16, bf16, or f32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1198,10 +1289,12 @@ LogicalResult VMIMulIOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(lhsType.getElementType())) + if (!isVMIIntegerLikeType(lhsType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + } + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) { return emitOpError("requires i8, i16, or i32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1210,8 +1303,9 @@ LogicalResult VMIFmaOp::verify() { auto rhsType = cast(getRhs().getType()); auto accType = cast(getAcc().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) { return emitOpError("requires f16, bf16, or f32 VMI element type"); + } return verifyFloatTernaryVRegOp(getOperation(), lhsType, rhsType, accType, resultType); } @@ -1224,10 +1318,12 @@ LogicalResult VMIDivFOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(lhsType.getElementType())) + if (!isVMIFloatLikeType(lhsType.getElementType())) { return emitOpError("requires floating-point-like VMI element type"); - if (!isVMIF16OrF32Type(lhsType.getElementType())) + } + if (!isVMIF16OrF32Type(lhsType.getElementType())) { return emitOpError("requires f16 or f32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1235,10 +1331,12 @@ LogicalResult VMIMinFOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(lhsType.getElementType())) + if (!isVMIFloatLikeType(lhsType.getElementType())) { return emitOpError("requires floating-point-like VMI element type"); - if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + } + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) { return emitOpError("requires f16, bf16, or f32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1246,10 +1344,12 @@ LogicalResult VMIMaxFOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(lhsType.getElementType())) + if (!isVMIFloatLikeType(lhsType.getElementType())) { return emitOpError("requires floating-point-like VMI element type"); - if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + } + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) { return emitOpError("requires f16, bf16, or f32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1276,24 +1376,27 @@ LogicalResult VMIMaxIOp::verify() { LogicalResult VMINegFOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIF16OrF32Type(sourceType.getElementType())) + if (!isVMIF16OrF32Type(sourceType.getElementType())) { return emitOpError("requires f16 or f32 VMI element type"); + } return verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType); } LogicalResult VMIAbsFOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIF16OrF32Type(sourceType.getElementType())) + if (!isVMIF16OrF32Type(sourceType.getElementType())) { return emitOpError("requires f16 or f32 VMI element type"); + } return verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType); } LogicalResult VMIAbsIOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(sourceType.getElementType())) + if (!isVMIIntegerLikeType(sourceType.getElementType())) { return emitOpError("requires integer-like VMI element type"); + } if (!isVMISignedOrSignlessI8I16I32Type(sourceType.getElementType())) return emitOpError("requires signless or signed i8, i16, or i32 VMI " "element type"); @@ -1305,32 +1408,36 @@ LogicalResult VMIAbsIOp::verify() { LogicalResult VMISqrtOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIF16OrF32Type(sourceType.getElementType())) + if (!isVMIF16OrF32Type(sourceType.getElementType())) { return emitOpError("requires f16 or f32 VMI element type"); + } return verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType); } LogicalResult VMIExpOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIF16OrF32Type(sourceType.getElementType())) + if (!isVMIF16OrF32Type(sourceType.getElementType())) { return emitOpError("requires f16 or f32 VMI element type"); + } return verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType); } LogicalResult VMILnOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIF16OrF32Type(sourceType.getElementType())) + if (!isVMIF16OrF32Type(sourceType.getElementType())) { return emitOpError("requires f16 or f32 VMI element type"); + } return verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType); } LogicalResult VMIReluOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIF16OrF32Type(sourceType.getElementType())) + if (!isVMIF16OrF32Type(sourceType.getElementType())) { return emitOpError("requires f16 or f32 VMI element type"); + } return verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType); } @@ -1338,10 +1445,12 @@ LogicalResult VMIAndIOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(lhsType.getElementType())) + if (!isVMIIntegerLikeType(lhsType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + } + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) { return emitOpError("requires i8, i16, or i32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1349,10 +1458,12 @@ LogicalResult VMIOrIOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(lhsType.getElementType())) + if (!isVMIIntegerLikeType(lhsType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + } + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) { return emitOpError("requires i8, i16, or i32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1360,10 +1471,12 @@ LogicalResult VMIXOrIOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(lhsType.getElementType())) + if (!isVMIIntegerLikeType(lhsType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + } + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) { return emitOpError("requires i8, i16, or i32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1371,10 +1484,12 @@ LogicalResult VMIShLIOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(lhsType.getElementType())) + if (!isVMIIntegerLikeType(lhsType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + } + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) { return emitOpError("requires i8, i16, or i32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1386,8 +1501,9 @@ LogicalResult VMIShRUIOp::verify() { if (!integerType || integerType.isSigned()) return emitOpError( "requires signless or unsigned integer VMI element type"); - if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) { return emitOpError("requires i8, i16, or i32 VMI element type"); + } return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); } @@ -1404,10 +1520,12 @@ LogicalResult VMIShRSIOp::verify() { LogicalResult VMINotOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(sourceType.getElementType())) + if (!isVMIIntegerLikeType(sourceType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (!isVMIAnyI8I16I32Type(sourceType.getElementType())) + } + if (!isVMIAnyI8I16I32Type(sourceType.getElementType())) { return emitOpError("requires i8, i16, or i32 VMI element type"); + } return verifyAllSameVRegShapeAndLayout(getOperation(), {sourceType, resultType}, /*requireSameElement=*/true); @@ -1419,10 +1537,12 @@ LogicalResult VMICmpFOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(lhsType.getElementType())) + if (!isVMIFloatLikeType(lhsType.getElementType())) { return emitOpError("requires floating-point-like VMI element type"); - if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + } + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) { return emitOpError("requires f16, bf16, or f32 VMI element type"); + } if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), {lhsType, rhsType}, /*requireSameElement=*/true))) return failure(); @@ -1433,10 +1553,12 @@ LogicalResult VMICmpIOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(lhsType.getElementType())) + if (!isVMIIntegerLikeType(lhsType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + } + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) { return emitOpError("requires i8, i16, or i32 VMI element type"); + } if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), {lhsType, rhsType}, /*requireSameElement=*/true))) return failure(); @@ -1459,11 +1581,13 @@ LogicalResult VMIActivePrefixIndexOp::verify() { auto maskType = cast(getMask().getType()); auto resultType = cast(getResult().getType()); auto resultIntType = dyn_cast(resultType.getElementType()); - if (!resultIntType || !resultIntType.isSignless()) + if (!resultIntType || !resultIntType.isSignless()) { return emitOpError("requires signless integer result element type"); + } unsigned resultWidth = resultIntType.getWidth(); - if (resultWidth != 8 && resultWidth != 16 && resultWidth != 32) + if (resultWidth != 8 && resultWidth != 16 && resultWidth != 32) { return emitOpError("requires i8, i16, or i32 result element type"); + } return verifyMaskMatchesData(getOperation(), maskType, resultType); } @@ -1501,15 +1625,19 @@ LogicalResult VMIReduceAddIOp::verify() { auto sourceType = cast(getSource().getType()); auto maskType = cast(getMask().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(sourceType.getElementType())) + if (!isVMIIntegerLikeType(sourceType.getElementType())) { return emitOpError("requires integer-like VMI source element type"); + } auto sourceIntegerType = dyn_cast(sourceType.getElementType()); - if (!sourceIntegerType || sourceIntegerType.getWidth() != 32) + if (!sourceIntegerType || sourceIntegerType.getWidth() != 32) { return emitOpError("requires 32-bit integer source element type"); - if (sourceType.getElementType() != resultType.getElementType()) + } + if (sourceType.getElementType() != resultType.getElementType()) { return emitOpError("requires source and result element types to match"); - if (resultType.getElementCount() != 1) + } + if (resultType.getElementCount() != 1) { return emitOpError("requires result to be a 1-lane VMI vector"); + } return verifyMaskMatchesData(getOperation(), maskType, sourceType); } @@ -1521,14 +1649,18 @@ LogicalResult VMIReduceAddFOp::verify() { return emitOpError( "requires reassoc attr because VPTO vcadd performs pair-wise " "floating-point reduction"); - if (!isVMIFloatLikeType(sourceType.getElementType())) + if (!isVMIFloatLikeType(sourceType.getElementType())) { return emitOpError("requires floating-point-like VMI source element type"); - if (!isVMIF16OrF32Type(sourceType.getElementType())) + } + if (!isVMIF16OrF32Type(sourceType.getElementType())) { return emitOpError("requires f16 or f32 source element type"); - if (sourceType.getElementType() != resultType.getElementType()) + } + if (sourceType.getElementType() != resultType.getElementType()) { return emitOpError("requires source and result element types to match"); - if (resultType.getElementCount() != 1) + } + if (resultType.getElementCount() != 1) { return emitOpError("requires result to be a 1-lane VMI vector"); + } return verifyMaskMatchesData(getOperation(), maskType, sourceType); } @@ -1539,12 +1671,15 @@ template LogicalResult verifyReduceMinMaxFOp(OpTy op) { if (!isVMIFloatLikeType(sourceType.getElementType())) return op.emitOpError( "requires floating-point-like VMI source element type"); - if (!isVMIF16OrF32Type(sourceType.getElementType())) + if (!isVMIF16OrF32Type(sourceType.getElementType())) { return op.emitOpError("requires f16 or f32 source element type"); - if (sourceType.getElementType() != resultType.getElementType()) + } + if (sourceType.getElementType() != resultType.getElementType()) { return op.emitOpError("requires source and result element types to match"); - if (resultType.getElementCount() != 1) + } + if (resultType.getElementCount() != 1) { return op.emitOpError("requires result to be a 1-lane VMI vector"); + } return verifyMaskMatchesData(op.getOperation(), maskType, sourceType); } @@ -1561,10 +1696,12 @@ template LogicalResult verifyReduceMinMaxIOp(OpTy op) { !isVMIAnyI8I16I32Type(sourceType.getElementType())) return op.emitOpError( "requires 8-bit, 16-bit, or 32-bit integer source element type"); - if (sourceType.getElementType() != resultType.getElementType()) + if (sourceType.getElementType() != resultType.getElementType()) { return op.emitOpError("requires source and result element types to match"); - if (resultType.getElementCount() != 1) + } + if (resultType.getElementCount() != 1) { return op.emitOpError("requires result to be a 1-lane VMI vector"); + } return verifyMaskMatchesData(op.getOperation(), maskType, sourceType); } @@ -1584,13 +1721,15 @@ static LogicalResult verifyGroupReduceFloatOp(OpTy op, bool requiresReassoc) { if (!isVMIFloatLikeType(sourceType.getElementType())) return op.emitOpError( "requires floating-point-like VMI source element type"); - if (!isVMIF16OrF32Type(sourceType.getElementType())) + if (!isVMIF16OrF32Type(sourceType.getElementType())) { return op.emitOpError("requires f16 or f32 source element type"); + } if (resultType.getElementCount() != op.getNumGroupsAttr().getInt()) return op.emitOpError( "requires result logical lane count to match num_groups"); - if (sourceType.getElementType() != resultType.getElementType()) + if (sourceType.getElementType() != resultType.getElementType()) { return op.emitOpError("requires source and result element types to match"); + } if (auto sourceLayout = sourceType.getLayoutAttr()) { bool supportedSourceLayout = sourceLayout.isContiguous() || @@ -1608,8 +1747,9 @@ static LogicalResult verifyGroupReduceFloatOp(OpTy op, bool requiresReassoc) { "#pto.vmi.layout"; } - if (failed(verifyMaskMatchesData(op.getOperation(), maskType, sourceType))) + if (failed(verifyMaskMatchesData(op.getOperation(), maskType, sourceType))) { return failure(); + } return verifyNumGroups(op.getOperation(), sourceType, op.getNumGroupsAttr().getInt()); } @@ -1631,8 +1771,9 @@ static LogicalResult verifyGroupReduceIntegerOp(OpTy op) { auto sourceType = cast(op.getSource().getType()); auto maskType = cast(op.getMask().getType()); auto resultType = cast(op.getResult().getType()); - if (!isVMIIntegerLikeType(sourceType.getElementType())) + if (!isVMIIntegerLikeType(sourceType.getElementType())) { return op.emitOpError("requires integer-like VMI source element type"); + } auto intType = dyn_cast(sourceType.getElementType()); if (!intType || !isVMIAnyI8I16I32Type(sourceType.getElementType())) return op.emitOpError( @@ -1640,8 +1781,9 @@ static LogicalResult verifyGroupReduceIntegerOp(OpTy op) { if (resultType.getElementCount() != op.getNumGroupsAttr().getInt()) return op.emitOpError( "requires result logical lane count to match num_groups"); - if (sourceType.getElementType() != resultType.getElementType()) + if (sourceType.getElementType() != resultType.getElementType()) { return op.emitOpError("requires source and result element types to match"); + } if (auto sourceLayout = sourceType.getLayoutAttr()) { bool supportedSourceLayout = sourceLayout.isContiguous() || @@ -1659,8 +1801,9 @@ static LogicalResult verifyGroupReduceIntegerOp(OpTy op) { "#pto.vmi.layout"; } - if (failed(verifyMaskMatchesData(op.getOperation(), maskType, sourceType))) + if (failed(verifyMaskMatchesData(op.getOperation(), maskType, sourceType))) { return failure(); + } return verifyNumGroups(op.getOperation(), sourceType, op.getNumGroupsAttr().getInt()); } @@ -1691,8 +1834,9 @@ LogicalResult VMIGroupBroadcastOp::verify() { if (resultType.getElementCount() % numGroups != 0) return emitOpError( "requires num_groups to evenly divide result logical lane count"); - if (sourceType.getElementType() != resultType.getElementType()) + if (sourceType.getElementType() != resultType.getElementType()) { return emitOpError("requires source and result element types to match"); + } if (auto sourceLayout = sourceType.getLayoutAttr()) { if (!sourceLayout.isGroupSlots() || sourceLayout.getNumGroups() != numGroups) @@ -1726,8 +1870,9 @@ template static LogicalResult verifyVMIHistogramOp(OpTy op) { if (resultType.getElementCount() != bins) return op.emitOpError("requires result element count to match acc " "(bins must be identical)"); - if (resultType.getLayoutAttr() != accType.getLayoutAttr()) + if (resultType.getLayoutAttr() != accType.getLayoutAttr()) { return op.emitOpError("requires result layout attribute to match acc"); + } auto resultElemType = dyn_cast(resultType.getElementType()); if (!resultElemType || resultElemType.getWidth() != 16 || !matchesVMIIntSemantics(resultElemType, VMIIntSignSemantics::Unsigned)) @@ -1737,8 +1882,9 @@ template static LogicalResult verifyVMIHistogramOp(OpTy op) { !matchesVMIIntSemantics(sourceElemType, VMIIntSignSemantics::Unsigned)) return op.emitOpError("requires source type to be " "!pto.vmi.vreg (interpreted as unsigned)"); - if (maskType.getElementCount() != sourceType.getElementCount()) + if (maskType.getElementCount() != sourceType.getElementCount()) { return op.emitOpError("requires mask logical lane count to match source"); + } if (auto accLayout = accType.getLayoutAttr()) { if (!accLayout.isContiguous()) @@ -1759,8 +1905,9 @@ template static LogicalResult verifyVMIHistogramOp(OpTy op) { if (!maskLayout.isContiguous()) return op.emitOpError("requires layout-assigned mask to use contiguous " "layout"); - if (maskType.getGranularity() != "b8") + if (maskType.getGranularity() != "b8") { return op.emitOpError("requires layout-assigned mask granularity b8"); + } } return success(); } @@ -1821,11 +1968,13 @@ LogicalResult VMITruncFOp::verify() { return emitOpError("rounding attr must be R, A, H, or Z"); } auto satAttr = (*this)->getAttrOfType("saturate"); - if (!satAttr) + if (!satAttr) { return emitOpError("'saturate' attribute is required (SAT or NOSAT)"); + } StringRef satVal = satAttr.getValue(); - if (satVal != "SAT" && satVal != "NOSAT") + if (satVal != "SAT" && satVal != "NOSAT") { return emitOpError("saturate attr must be 'SAT' or 'NOSAT'"); + } return success(); } @@ -1835,24 +1984,27 @@ LogicalResult VMIFPToSIOp::verify() { if (sourceType.getElementCount() != resultType.getElementCount()) return emitOpError( "requires source and result logical lane counts to match"); - if (!isVMIFloatLikeType(sourceType.getElementType())) + if (!isVMIFloatLikeType(sourceType.getElementType())) { return emitOpError("requires floating-point-like source element type"); + } if (!isVMISignedOrSignlessIntegerType(resultType.getElementType())) return emitOpError("requires signed or signless integer result element " "type"); auto contract = lookupVMIFpToSiContract(sourceType.getElementType(), resultType.getElementType()); - if (!contract) + if (!contract) { return emitOpError("unsupported fp-to-si conversion element type pair"); + } if (contract->requiresSat) { auto satAttr = (*this)->getAttrOfType("saturate"); if (!satAttr) return emitOpError("'saturate' attribute is required for this fp-to-si " "conversion (SAT or NOSAT)"); StringRef satVal = satAttr.getValue(); - if (satVal != "SAT" && satVal != "NOSAT") + if (satVal != "SAT" && satVal != "NOSAT") { return emitOpError("saturate attr must be 'SAT' or 'NOSAT'"); + } } else { if ((*this)->getAttrOfType("saturate")) return emitOpError("'saturate' attribute is not valid for this fp-to-si " @@ -1867,23 +2019,27 @@ LogicalResult VMIFPToUIOp::verify() { if (sourceType.getElementCount() != resultType.getElementCount()) return emitOpError( "requires source and result logical lane counts to match"); - if (!isVMIFloatLikeType(sourceType.getElementType())) + if (!isVMIFloatLikeType(sourceType.getElementType())) { return emitOpError("requires floating-point-like source element type"); - if (!isVMIUnsignedIntegerType(resultType.getElementType())) + } + if (!isVMIUnsignedIntegerType(resultType.getElementType())) { return emitOpError("requires unsigned integer result element type"); + } auto contract = lookupVMIFpToUIContract(sourceType.getElementType(), resultType.getElementType()); - if (!contract) + if (!contract) { return emitOpError("unsupported fp-to-ui conversion element type pair"); + } if (contract->requiresSat) { auto satAttr = (*this)->getAttrOfType("saturate"); if (!satAttr) return emitOpError("'saturate' attribute is required for this fp-to-ui " "conversion (SAT or NOSAT)"); StringRef satVal = satAttr.getValue(); - if (satVal != "SAT" && satVal != "NOSAT") + if (satVal != "SAT" && satVal != "NOSAT") { return emitOpError("saturate attr must be 'SAT' or 'NOSAT'"); + } } else { if ((*this)->getAttrOfType("saturate")) return emitOpError("'saturate' attribute is not valid for this fp-to-ui " @@ -1901,12 +2057,15 @@ LogicalResult VMISIToFPOp::verify() { if (!isVMISignedOrSignlessIntegerType(sourceType.getElementType())) return emitOpError( "requires signed or signless integer source element type"); - if (!isVMIFloatLikeType(resultType.getElementType())) + if (!isVMIFloatLikeType(resultType.getElementType())) { return emitOpError("requires floating-point-like result element type"); - if (getVMIElementBitWidth(sourceType.getElementType()) != 32) + } + if (getVMIElementBitWidth(sourceType.getElementType()) != 32) { return emitOpError("requires 32-bit integer source element type"); - if (!resultType.getElementType().isF32()) + } + if (!resultType.getElementType().isF32()) { return emitOpError("requires f32 result element type"); + } return success(); } @@ -1958,11 +2117,13 @@ LogicalResult VMITruncIOp::verify() { return emitOpError( "requires result element type to be narrower than source element type"); auto satAttr = (*this)->getAttrOfType("saturate"); - if (!satAttr) + if (!satAttr) { return emitOpError("'saturate' attribute is required (SAT or NOSAT)"); + } StringRef satVal = satAttr.getValue(); - if (satVal != "SAT" && satVal != "NOSAT") + if (satVal != "SAT" && satVal != "NOSAT") { return emitOpError("saturate attr must be 'SAT' or 'NOSAT'"); + } return success(); } @@ -1986,8 +2147,9 @@ LogicalResult VMIBitcastOp::verify() { return emitOpError( "requires either both source and result to carry layout or neither " "to carry layout"); - if (sourceType.getLayout() != resultType.getLayout()) + if (sourceType.getLayout() != resultType.getLayout()) { return emitOpError("requires source and result layouts to match"); + } } return success(); @@ -2083,8 +2245,9 @@ void VMIGroupSlotLoadOp::getEffects( LogicalResult VMIGroupBroadcastLoadOp::verify() { auto resultType = cast(getResult().getType()); int64_t numGroups = getNumGroupsAttr().getInt(); - if (numGroups <= 0) + if (numGroups <= 0) { return emitOpError("requires num_groups to be positive"); + } if (resultType.getElementCount() % numGroups != 0) return emitOpError( "requires num_groups to evenly divide result logical lane count"); @@ -2374,8 +2537,9 @@ LogicalResult VMIShuffleOp::verify() { LogicalResult VMIChannelSplitOp::verify() { auto sourceType = cast(getSource().getType()); - if (getResults().size() < 2) + if (getResults().size() < 2) { return emitOpError("requires at least two channel results"); + } auto firstResultType = cast(getResults().front().getType()); if (sourceType.getElementCount() != static_cast(getResults().size()) * @@ -2390,8 +2554,9 @@ LogicalResult VMIChannelSplitOp::verify() { "count and source element type"); } bool anyLayout = isLayoutAssigned(sourceType); - for (Value result : getResults()) + for (Value result : getResults()) { anyLayout |= isLayoutAssigned(cast(result.getType())); + } if (anyLayout) { if (!isLayoutAssigned(sourceType)) return emitOpError("requires layout-assigned channel_split source when " @@ -2419,8 +2584,9 @@ LogicalResult VMIChannelSplitOp::verify() { } LogicalResult VMIChannelMergeOp::verify() { - if (getInputs().size() < 2) + if (getInputs().size() < 2) { return emitOpError("requires at least two channel inputs"); + } auto firstInputType = cast(getInputs().front().getType()); auto resultType = cast(getResult().getType()); for (Value input : getInputs()) { @@ -2436,8 +2602,9 @@ LogicalResult VMIChannelMergeOp::verify() { return emitOpError( "requires result lane count and element type to match merged channels"); bool anyLayout = isLayoutAssigned(resultType); - for (Value input : getInputs()) + for (Value input : getInputs()) { anyLayout |= isLayoutAssigned(cast(input.getType())); + } if (anyLayout) { if (!isLayoutAssigned(resultType)) return emitOpError("requires layout-assigned channel_merge result when " @@ -2471,8 +2638,9 @@ LogicalResult VMIEnsureLayoutOp::verify() { sourceType.getElementType() != resultType.getElementType()) return emitOpError("requires source and result to preserve VMI data shape " "and element type"); - if (!isLayoutAssigned(sourceType) || !isLayoutAssigned(resultType)) + if (!isLayoutAssigned(sourceType) || !isLayoutAssigned(resultType)) { return emitOpError("requires source and result to be layout-assigned"); + } return success(); } @@ -2483,8 +2651,9 @@ LogicalResult VMIEnsureMaskLayoutOp::verify() { sourceType.getGranularity() != resultType.getGranularity()) return emitOpError("requires source and result to preserve VMI mask shape " "and granularity"); - if (!isLayoutAssigned(sourceType) || !isLayoutAssigned(resultType)) + if (!isLayoutAssigned(sourceType) || !isLayoutAssigned(resultType)) { return emitOpError("requires source and result to be layout-assigned"); + } return success(); } @@ -2523,8 +2692,9 @@ enum class CvtDirection { FpWiden, FpNarrow, FpToSi, FpToUi, SiToFp, IntWiden, I static LogicalResult verifyVMIPmodeMask(Operation *op, VMIMaskType maskType, VMIVRegType dataType, std::optional pmode) { - if (failed(verifyMaskMatchesData(op, maskType, dataType))) + if (failed(verifyMaskMatchesData(op, maskType, dataType))) { return failure(); + } if (pmode.has_value()) { StringRef mode = pmode.value(); if (mode != "merge" && mode != "zero") @@ -2545,10 +2715,12 @@ static LogicalResult verifyVMIVariadicPmodeMask(Operation *op, return op->emitOpError("pmode must be \"merge\" or \"zero\", got \"") << mode << "\""; } - if (maskParts.empty()) + if (maskParts.empty()) { return success(); - if (maskParts.size() != 1) + } + if (maskParts.size() != 1) { return op->emitOpError("expects at most one mask operand"); + } return verifyMaskMatchesData(op, cast(maskParts.front().getType()), dataType); } @@ -2577,8 +2749,9 @@ verifyVMIVectorScalarOp(Operation *op, VMIVRegType srcType, /*requireSameElement=*/true))) return failure(); - if (failed(verifyMaskMatchesData(op, maskType, resultType))) + if (failed(verifyMaskMatchesData(op, maskType, resultType))) { return failure(); + } if (pmode.has_value()) { StringRef mode = pmode.value(); @@ -2600,13 +2773,15 @@ verifyVMIVectorScalarShiftOp(Operation *op, VMIVRegType srcType, if (!isVMIIntegerLikeType(eltTy)) return op->emitOpError( "requires integer-like VMI element type for shift"); - if (!scalarType.isSignlessInteger(16)) + if (!scalarType.isSignlessInteger(16)) { return op->emitOpError("requires signless i16 shift amount"); + } if (failed(verifyAllSameVRegShapeAndLayout(op, {srcType, resultType}, /*requireSameElement=*/true))) return failure(); - if (failed(verifyMaskMatchesData(op, maskType, resultType))) + if (failed(verifyMaskMatchesData(op, maskType, resultType))) { return failure(); + } if (pmode.has_value()) { StringRef mode = pmode.value(); if (mode != "merge" && mode != "zero") @@ -2695,14 +2870,16 @@ LogicalResult VMIVbrcOp::verify() { // Group broadcast mode int64_t numGroupsVal = groupAttr.getInt(); auto vregType = dyn_cast(valueType); - if (!vregType) + if (!vregType) { return emitOpError("requires VMI vector input when num_groups is set"); + } if (vregType.getElementCount() != numGroupsVal) return emitOpError() << "requires source logical lane count " << vregType.getElementCount() << " to match num_groups " << numGroupsVal; - if (vregType.getElementType() != resultType.getElementType()) + if (vregType.getElementType() != resultType.getElementType()) { return emitOpError("requires source and result element types to match"); + } if (auto sourceLayout = vregType.getLayoutAttr()) { if (!sourceLayout.isGroupSlots() || sourceLayout.getNumGroups() != numGroupsVal) @@ -2719,11 +2896,13 @@ LogicalResult VMIVbrcOp::verify() { } // Scalar/1-lane broadcast mode (no num_groups) - if (valueType == resultType.getElementType()) + if (valueType == resultType.getElementType()) { return success(); + } if (auto vregType = dyn_cast(valueType)) { - if (vregType.getElementCount() != 1) + if (vregType.getElementCount() != 1) { return emitOpError("requires VMI vector input to have one logical lane"); + } if (vregType.getElementType() != resultType.getElementType()) return emitOpError("requires VMI vector input element type to match " "result element type"); @@ -2739,17 +2918,20 @@ LogicalResult VMIVciOp::verify() { if (!isVMIIotaElementType(elementType)) return emitOpError("requires result element type to be integer 8/16/32 " "or f16/f32"); - if (!isCompatibleScalarForSemanticType(elementType, getBase().getType())) + if (!isCompatibleScalarForSemanticType(elementType, getBase().getType())) { return emitOpError("requires base type to match result element type"); + } if (std::optional order = getOrder()) { - if (*order != "ASC" && *order != "DESC") + if (*order != "ASC" && *order != "DESC") { return emitOpError("requires order to be ASC or DESC"); + } } if (auto groupAttr = getGroupAttr()) { int64_t numGroups = groupAttr.getInt(); - if (numGroups <= 0) + if (numGroups <= 0) { return emitOpError("requires group to be positive"); + } if (resultType.getElementCount() % numGroups != 0) return emitOpError("requires group to evenly divide result logical lane " "count"); @@ -2769,41 +2951,50 @@ LogicalResult VMIVciOp::verify() { LogicalResult VMIPsetOp::verify() { auto resultType = cast(getResult().getType()); StringRef pattern = getPattern(); - if (pattern != "PAT_ALL") + if (pattern != "PAT_ALL") { return emitOpError("requires pattern to be \"PAT_ALL\""); - if (!resultType.isPred() && !isLayoutAssigned(resultType)) + } + if (!resultType.isPred() && !isLayoutAssigned(resultType)) { return emitOpError("requires concrete mask result to carry layout"); + } return success(); } LogicalResult VMIPgeOp::verify() { auto resultType = cast(getResult().getType()); StringRef pattern = getPattern(); - if (!pattern.starts_with("PAT_VL")) + if (!pattern.starts_with("PAT_VL")) { return emitOpError("requires pattern to start with \"PAT_VL\""); + } int64_t activeLanes; - if (pattern.drop_front(6).getAsInteger(10, activeLanes)) + if (pattern.drop_front(6).getAsInteger(10, activeLanes)) { return emitOpError("requires pattern \"PAT_VL\" with integer n"); - if (activeLanes <= 0) + } + if (activeLanes <= 0) { return emitOpError("requires positive n in pattern \"PAT_VL\""); + } if (activeLanes > resultType.getElementCount()) return emitOpError("PAT_VL active lanes ") << activeLanes << " exceeds mask element count " << resultType.getElementCount(); - if (!resultType.isPred() && !isLayoutAssigned(resultType)) + if (!resultType.isPred() && !isLayoutAssigned(resultType)) { return emitOpError("requires concrete mask result to carry layout"); + } return success(); } LogicalResult VMIPltOp::verify() { auto resultType = cast(getMask().getType()); auto scalarType = dyn_cast(getScalar().getType()); - if (!scalarType || scalarType.getWidth() != 32) + if (!scalarType || scalarType.getWidth() != 32) { return emitOpError("requires i32 scalar input"); + } auto scalarOutType = dyn_cast(getScalarOut().getType()); - if (!scalarOutType || scalarOutType.getWidth() != 32) + if (!scalarOutType || scalarOutType.getWidth() != 32) { return emitOpError("requires i32 scalar_out result"); - if (!resultType.isPred() && !isLayoutAssigned(resultType)) + } + if (!resultType.isPred() && !isLayoutAssigned(resultType)) { return emitOpError("requires concrete mask result to carry layout"); + } return success(); } @@ -2811,8 +3002,9 @@ LogicalResult VMIVaddOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -2823,8 +3015,9 @@ LogicalResult VMIVsubOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -2835,8 +3028,9 @@ LogicalResult VMIVmulOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -2847,10 +3041,12 @@ LogicalResult VMIVdivOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(lhsType.getElementType())) + if (!isVMIFloatLikeType(lhsType.getElementType())) { return emitOpError("requires floating-point-like VMI element type"); - if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + } + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -2894,8 +3090,9 @@ LogicalResult VMIVmaxOp::verify() { LogicalResult VMIVnegOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) + if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -2925,8 +3122,9 @@ LogicalResult VMIVabsOp::verify() { LogicalResult VMIVsqrtOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) + if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -2936,8 +3134,9 @@ LogicalResult VMIVsqrtOp::verify() { LogicalResult VMIVexpOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) + if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -2947,8 +3146,9 @@ LogicalResult VMIVexpOp::verify() { LogicalResult VMIVlnOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) + if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -2958,8 +3158,9 @@ LogicalResult VMIVlnOp::verify() { LogicalResult VMIVreluOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) + if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -2969,10 +3170,12 @@ LogicalResult VMIVreluOp::verify() { LogicalResult VMIVandOp::verify() { if (isa(getLhs().getType())) { // Mask logic path: reject predication mask and pmode. - if (!getMask().empty()) + if (!getMask().empty()) { return emitOpError("mask logic op does not support predication mask"); - if (auto pmode = getPmode()) + } + if (auto pmode = getPmode()) { return emitOpError("mask logic op does not support pmode"); + } auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); @@ -2983,10 +3186,12 @@ LogicalResult VMIVandOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(lhsType.getElementType())) + if (!isVMIIntegerLikeType(lhsType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + } + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -2996,10 +3201,12 @@ LogicalResult VMIVandOp::verify() { LogicalResult VMIVorOp::verify() { if (isa(getLhs().getType())) { // Mask logic path: reject predication mask and pmode. - if (!getMask().empty()) + if (!getMask().empty()) { return emitOpError("mask logic op does not support predication mask"); - if (auto pmode = getPmode()) + } + if (auto pmode = getPmode()) { return emitOpError("mask logic op does not support pmode"); + } auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); @@ -3010,10 +3217,12 @@ LogicalResult VMIVorOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(lhsType.getElementType())) + if (!isVMIIntegerLikeType(lhsType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + } + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -3023,10 +3232,12 @@ LogicalResult VMIVorOp::verify() { LogicalResult VMIVxorOp::verify() { if (isa(getLhs().getType())) { // Mask logic path: reject predication mask and pmode. - if (!getMask().empty()) + if (!getMask().empty()) { return emitOpError("mask logic op does not support predication mask"); - if (auto pmode = getPmode()) + } + if (auto pmode = getPmode()) { return emitOpError("mask logic op does not support pmode"); + } auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); @@ -3037,10 +3248,12 @@ LogicalResult VMIVxorOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(lhsType.getElementType())) + if (!isVMIIntegerLikeType(lhsType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + } + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -3051,10 +3264,12 @@ LogicalResult VMIVshlOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(lhsType.getElementType())) + if (!isVMIIntegerLikeType(lhsType.getElementType())) { return emitOpError("requires integer-like VMI element type"); - if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + } + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -3066,10 +3281,12 @@ LogicalResult VMIVshrOp::verify() { auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); auto integerType = dyn_cast(lhsType.getElementType()); - if (!integerType) + if (!integerType) { return emitOpError("requires integer VMI element type"); - if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + } + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) { return failure(); + } if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), resultType, getPmode()))) return failure(); @@ -3079,10 +3296,12 @@ LogicalResult VMIVshrOp::verify() { LogicalResult VMIVnotOp::verify() { if (isa(getSource().getType())) { // Mask logic path: reject predication mask and pmode. - if (!getMask().empty()) + if (!getMask().empty()) { return emitOpError("mask logic op does not support predication mask"); - if (auto pmode = getPmode()) + } + if (auto pmode = getPmode()) { return emitOpError("mask logic op does not support pmode"); + } auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); return verifyAllSameMaskShapeLayoutAndGranularity( @@ -3091,8 +3310,9 @@ LogicalResult VMIVnotOp::verify() { // VReg path. auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIIntegerLikeType(sourceType.getElementType())) + if (!isVMIIntegerLikeType(sourceType.getElementType())) { return emitOpError("requires integer-like VMI element type"); + } if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), {sourceType, resultType}, /*requireSameElement=*/true))) @@ -3112,8 +3332,9 @@ LogicalResult VMIvSelOp::verify() { {trueType, falseType, resultType}, /*requireSameElement=*/true))) return failure(); - if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) { return failure(); + } if (auto pmode = getPmode(); pmode.has_value()) { StringRef mode = pmode.value(); if (mode != "merge" && mode != "zero") @@ -3137,11 +3358,13 @@ LogicalResult VMIvcaddOp::verify() { "source element type"); // Floating-point vcadd MUST carry reassoc - if (isFloat && !getReassoc()) + if (isFloat && !getReassoc()) { return emitOpError("floating add-reduction requires reassoc attr"); + } - if (failed(verifyMaskMatchesData(getOperation(), maskType, sourceType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, sourceType))) { return failure(); + } // Validate group vs result lane count if (auto groupAttr = getGroupAttr()) { @@ -3167,14 +3390,16 @@ LogicalResult VMIvcaddOp::verify() { } // Element types must match - if (sourceType.getElementType() != resultType.getElementType()) + if (sourceType.getElementType() != resultType.getElementType()) { return emitOpError("source and result element types must match"); + } // pmode must be "zero" or "merge" if set if (auto pmode = getPmode()) { StringRef val = *pmode; - if (val != "zero" && val != "merge") + if (val != "zero" && val != "merge") { return emitOpError("pmode must be \"zero\" or \"merge\", got \"") << val << "\""; + } } return success(); @@ -3192,8 +3417,9 @@ LogicalResult VMIvcmaxOp::verify() { return emitOpError("requires integer-like or floating-point-like VMI " "source element type"); - if (failed(verifyMaskMatchesData(getOperation(), maskType, sourceType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, sourceType))) { return failure(); + } if (auto groupAttr = getGroupAttr()) { int64_t C = groupAttr.getInt(); @@ -3217,13 +3443,15 @@ LogicalResult VMIvcmaxOp::verify() { << resultType.getElementCount(); } - if (sourceType.getElementType() != resultType.getElementType()) + if (sourceType.getElementType() != resultType.getElementType()) { return emitOpError("source and result element types must match"); + } if (auto pmode = getPmode()) { StringRef val = *pmode; - if (val != "zero" && val != "merge") + if (val != "zero" && val != "merge") { return emitOpError("pmode must be \"zero\" or \"merge\", got \"") << val << "\""; + } } return success(); @@ -3241,8 +3469,9 @@ LogicalResult VMIvcminOp::verify() { return emitOpError("requires integer-like or floating-point-like VMI " "source element type"); - if (failed(verifyMaskMatchesData(getOperation(), maskType, sourceType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, sourceType))) { return failure(); + } if (auto groupAttr = getGroupAttr()) { int64_t C = groupAttr.getInt(); @@ -3266,13 +3495,15 @@ LogicalResult VMIvcminOp::verify() { << resultType.getElementCount(); } - if (sourceType.getElementType() != resultType.getElementType()) + if (sourceType.getElementType() != resultType.getElementType()) { return emitOpError("source and result element types must match"); + } if (auto pmode = getPmode()) { StringRef val = *pmode; - if (val != "zero" && val != "merge") + if (val != "zero" && val != "merge") { return emitOpError("pmode must be \"zero\" or \"merge\", got \"") << val << "\""; + } } return success(); @@ -3302,8 +3533,9 @@ LogicalResult VMIVgatherOp::verify() { getOperation(), {offsetsType, resultType}, /*requireSameElement=*/false))) return failure(); - if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) { return failure(); + } // 16-bit offsets only address the ui16 gather path (pto.vgather2 / b16 mask), // which requires a ui16 result element type. Reject other 16-bit-offset @@ -3317,8 +3549,9 @@ LogicalResult VMIVgatherOp::verify() { "requires ui16 result element type when using ui16 offsets"); if (auto pmode = getPmode()) { - if (pmode.value() != "merge" && pmode.value() != "zero") + if (pmode.value() != "merge" && pmode.value() != "zero") { return emitOpError("pmode must be 'merge' or 'zero'"); + } } return success(); } @@ -3349,12 +3582,14 @@ LogicalResult VMIVgatherbOp::verify() { return emitOpError( "requires signless or unsigned 16-bit or 32-bit integer offsets"); - if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) { return failure(); + } if (auto pmode = getPmode()) { - if (pmode.value() != "merge" && pmode.value() != "zero") + if (pmode.value() != "merge" && pmode.value() != "zero") { return emitOpError("pmode must be 'merge' or 'zero'"); + } } return success(); } @@ -3391,12 +3626,14 @@ LogicalResult VMIVscatterOp::verify() { {valueType, offsetsType}, /*requireSameElement=*/false))) return failure(); - if (failed(verifyMaskMatchesData(getOperation(), maskType, valueType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, valueType))) { return failure(); + } if (auto pmode = getPmode()) { - if (pmode.value() != "merge" && pmode.value() != "zero") + if (pmode.value() != "merge" && pmode.value() != "zero") { return emitOpError("pmode must be 'merge' or 'zero'"); + } } return success(); } @@ -3413,28 +3650,33 @@ LogicalResult VMIVexpdifOp::verify() { auto maskType = cast(getMask().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(xType.getElementType())) + if (!isVMIFloatLikeType(xType.getElementType())) { return emitOpError("requires x element type to be f16 or f32"); + } auto maxElemType = dyn_cast(maxType.getElementType()); - if (!maxElemType || maxElemType.getWidth() != 32) + if (!maxElemType || maxElemType.getWidth() != 32) { return emitOpError("requires max element type to be f32"); + } auto resultElemType = dyn_cast(resultType.getElementType()); - if (!resultElemType || resultElemType.getWidth() != 32) + if (!resultElemType || resultElemType.getWidth() != 32) { return emitOpError("requires result element type to be f32"); + } if (xType.getElementCount() != maxType.getElementCount() || xType.getElementCount() != resultType.getElementCount()) return emitOpError( "requires x, max, and result logical lane counts to match"); - if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) { return failure(); + } if (auto pmode = getPmode()) { - if (pmode.value() != "merge" && pmode.value() != "zero") + if (pmode.value() != "merge" && pmode.value() != "zero") { return emitOpError("pmode must be 'merge' or 'zero'"); + } } return success(); } @@ -3445,23 +3687,27 @@ LogicalResult VMIVaxpyOp::verify() { auto maskType = cast(getMask().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(xType.getElementType())) + if (!isVMIFloatLikeType(xType.getElementType())) { return emitOpError("requires vector element type to be f16 or f32"); + } if (xType != accType || accType != resultType) return emitOpError( "requires x, acc, and result to have identical VMI vreg types"); auto alphaType = cast(getAlpha().getType()); - if (alphaType != xType.getElementType()) + if (alphaType != xType.getElementType()) { return emitOpError("requires alpha scalar type to match vector element type"); + } - if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) { return failure(); + } if (auto pmode = getPmode()) { - if (pmode.value() != "merge" && pmode.value() != "zero") + if (pmode.value() != "merge" && pmode.value() != "zero") { return emitOpError("pmode must be 'merge' or 'zero'"); + } } return success(); } @@ -3471,23 +3717,27 @@ LogicalResult VMIVlreluOp::verify() { auto maskType = cast(getMask().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(xType.getElementType())) + if (!isVMIFloatLikeType(xType.getElementType())) { return emitOpError("requires vector element type to be f16 or f32"); + } - if (xType != resultType) + if (xType != resultType) { return emitOpError("requires x and result to have identical VMI vreg types"); + } auto slopeType = cast(getSlope().getType()); if (slopeType != xType.getElementType()) return emitOpError( "requires slope scalar type to match vector element type"); - if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) { return failure(); + } if (auto pmode = getPmode()) { - if (pmode.value() != "merge" && pmode.value() != "zero") + if (pmode.value() != "merge" && pmode.value() != "zero") { return emitOpError("pmode must be 'merge' or 'zero'"); + } } return success(); } @@ -3498,19 +3748,22 @@ LogicalResult VMIVpreluOp::verify() { auto maskType = cast(getMask().getType()); auto resultType = cast(getResult().getType()); - if (!isVMIFloatLikeType(xType.getElementType())) + if (!isVMIFloatLikeType(xType.getElementType())) { return emitOpError("requires vector element type to be f16 or f32"); + } if (xType != alphaType || alphaType != resultType) return emitOpError( "requires x, alpha, and result to have identical VMI vreg types"); - if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) { return failure(); + } if (auto pmode = getPmode()) { - if (pmode.value() != "merge" && pmode.value() != "zero") + if (pmode.value() != "merge" && pmode.value() != "zero") { return emitOpError("pmode must be 'merge' or 'zero'"); + } } return success(); } @@ -3578,14 +3831,17 @@ LogicalResult VMIVmullOp::verify() { "requires a, b, low, and high to have identical VMI vreg types"); int64_t lanes = aType.getElementCount(); - if (lanes != 64 && lanes != 128 && lanes != 256) + if (lanes != 64 && lanes != 128 && lanes != 256) { return emitOpError("requires logical lane count to be 64, 128, or 256"); + } - if (failed(verifyMaskMatchesData(getOperation(), maskType, aType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, aType))) { return failure(); + } - if (auto pmode = getPmode(); pmode && pmode.value() != "zero") + if (auto pmode = getPmode(); pmode && pmode.value() != "zero") { return emitOpError("pmode must be 'zero' when specified"); + } return success(); } @@ -3635,10 +3891,12 @@ LogicalResult VMICvtOp::verify() { // 2. Classify the conversion direction. CvtDirection dir; if (srcFp && dstFp) { - if (dstBits > srcBits) + if (dstBits > srcBits) { dir = CvtDirection::FpWiden; - else if (dstBits < srcBits) + } + else if (dstBits < srcBits) { dir = CvtDirection::FpNarrow; + } else { // Same-width fp→fp (e.g. bf16 → f16): only allowed for VMI fp-to-fp // contract pairs, routed through FpNarrow (1:1 TruncF). @@ -3649,10 +3907,12 @@ LogicalResult VMICvtOp::verify() { dir = CvtDirection::FpNarrow; } } else if (srcFp && dstInt) { - if (isVMIUnsignedIntegerType(dstElem)) + if (isVMIUnsignedIntegerType(dstElem)) { dir = CvtDirection::FpToUi; - else + } + else { dir = CvtDirection::FpToSi; + } } else if (srcInt && dstFp) { if (!isVMISignedOrSignlessIntegerType(srcElem)) return emitOpError( @@ -3660,10 +3920,12 @@ LogicalResult VMICvtOp::verify() { "element type"); dir = CvtDirection::SiToFp; } else if (srcInt && dstInt) { - if (dstBits > srcBits) + if (dstBits > srcBits) { dir = CvtDirection::IntWiden; - else if (dstBits < srcBits) + } + else if (dstBits < srcBits) { dir = CvtDirection::IntNarrow; + } else return emitOpError( "int-to-int conversion must change element bit-width"); @@ -3690,15 +3952,17 @@ LogicalResult VMICvtOp::verify() { auto satAttr = (*this)->getAttrOfType("saturate"); if (dir == CvtDirection::FpToSi) { auto contract = lookupVMIFpToSiContract(srcElem, dstElem); - if (!contract) + if (!contract) { return emitOpError("unsupported fp-to-si conversion element type pair"); + } if (contract->requiresSat) { if (!satAttr) return emitOpError("'saturate' attribute is required for this " "fp-to-si conversion; write 'SAT' or 'NOSAT'"); StringRef satVal = satAttr.getValue(); - if (satVal != "SAT" && satVal != "NOSAT") + if (satVal != "SAT" && satVal != "NOSAT") { return emitOpError("saturate must be 'SAT' or 'NOSAT'"); + } } else { if (satAttr) return emitOpError("'saturate' attribute is not valid for this " @@ -3706,15 +3970,17 @@ LogicalResult VMICvtOp::verify() { } } else if (dir == CvtDirection::FpToUi) { auto contract = lookupVMIFpToUIContract(srcElem, dstElem); - if (!contract) + if (!contract) { return emitOpError("unsupported fp-to-ui conversion element type pair"); + } if (contract->requiresSat) { if (!satAttr) return emitOpError("'saturate' attribute is required for this " "fp-to-ui conversion; write 'SAT' or 'NOSAT'"); StringRef satVal = satAttr.getValue(); - if (satVal != "SAT" && satVal != "NOSAT") + if (satVal != "SAT" && satVal != "NOSAT") { return emitOpError("saturate must be 'SAT' or 'NOSAT'"); + } } else { if (satAttr) return emitOpError("'saturate' attribute is not valid for this " @@ -3729,8 +3995,9 @@ LogicalResult VMICvtOp::verify() { "int-narrow conversions; write 'SAT' or " "'NOSAT'"); StringRef satVal = satAttr.getValue(); - if (satVal != "SAT" && satVal != "NOSAT") + if (satVal != "SAT" && satVal != "NOSAT") { return emitOpError("saturate must be 'SAT' or 'NOSAT'"); + } // si32 -> si8 IntNarrow has no native hardware form. Lowering aliases // it through ui32 -> ui8 (bit-pattern equal ONLY under NOSAT). Reject // SAT here because ui32 -> ui8 SAT clamps to [0, 255], which does NOT @@ -3763,8 +4030,9 @@ LogicalResult VMICvtOp::verify() { // --- pmode --- if (auto pmodeAttr = (*this)->getAttrOfType("pmode")) { StringRef pmode = pmodeAttr.getValue(); - if (pmode != "merge" && pmode != "zero") + if (pmode != "merge" && pmode != "zero") { return emitOpError("pmode must be 'merge' or 'zero'"); + } } return success(); @@ -3790,8 +4058,9 @@ LogicalResult VMIVinterpretCastOp::verify() { return emitOpError( "requires either both source and result to carry layout or neither " "to carry layout"); - if (sourceType.getLayout() != resultType.getLayout()) + if (sourceType.getLayout() != resultType.getLayout()) { return emitOpError("requires source and result layouts to match"); + } } return success(); @@ -3803,28 +4072,33 @@ ParseResult VMIvStoreOp::parse(OpAsmParser &parser, OperationState &result) { OpAsmParser::UnresolvedOperand offsetOperand; SmallVector postBracketOps; - if (parser.parseOperand(operand)) + if (parser.parseOperand(operand)) { return failure(); + } preBracketOperands.push_back(operand); bool consumedLSquare = false; while (!consumedLSquare) { if (succeeded(parser.parseOptionalLSquare())) { - if (parser.parseOperand(offsetOperand) || parser.parseRSquare()) + if (parser.parseOperand(offsetOperand) || parser.parseRSquare()) { return failure(); + } consumedLSquare = true; break; } - if (parser.parseComma()) + if (parser.parseComma()) { return failure(); + } if (succeeded(parser.parseOptionalLSquare())) { - if (parser.parseOperand(offsetOperand) || parser.parseRSquare()) + if (parser.parseOperand(offsetOperand) || parser.parseRSquare()) { return failure(); + } consumedLSquare = true; break; } - if (parser.parseOperand(operand)) + if (parser.parseOperand(operand)) { return failure(); + } preBracketOperands.push_back(operand); } @@ -3836,19 +4110,23 @@ ParseResult VMIvStoreOp::parse(OpAsmParser &parser, OperationState &result) { // Up to 3, disambiguated after parsing attrs. while (succeeded(parser.parseOptionalComma())) { OpAsmParser::UnresolvedOperand postOp; - if (parser.parseOperand(postOp)) + if (parser.parseOperand(postOp)) { return failure(); + } postBracketOps.push_back(postOp); - if (postBracketOps.size() >= 3) + if (postBracketOps.size() >= 3) { break; + } } - if (parser.parseOptionalAttrDict(result.attributes)) + if (parser.parseOptionalAttrDict(result.attributes)) { return failure(); + } SmallVector types; - if (parser.parseColon() || parser.parseTypeList(types)) + if (parser.parseColon() || parser.parseTypeList(types)) { return failure(); + } bool hasGroup = result.attributes.get("group") != nullptr; bool hasStride = false; @@ -3897,8 +4175,9 @@ ParseResult VMIvStoreOp::parse(OpAsmParser &parser, OperationState &result) { << "), got " << nTypes; for (size_t i = 0; i < nValues; ++i) { - if (parser.resolveOperand(preBracketOperands[i], types[i], result.operands)) + if (parser.resolveOperand(preBracketOperands[i], types[i], result.operands)) { return failure(); + } } Type destType = types[nValues]; @@ -3943,8 +4222,9 @@ ParseResult VMIvStoreOp::parse(OpAsmParser &parser, OperationState &result) { } void VMIvStoreOp::print(OpAsmPrinter &p) { - for (auto val : getValues()) + for (auto val : getValues()) { p << ' ' << val << ", "; + } p << getDestination() << '['; p.printOperand(getOffset()); p << ']'; @@ -3964,33 +4244,42 @@ void VMIvStoreOp::print(OpAsmPrinter &p) { } p.printOptionalAttrDict((*this)->getAttrs(), {"operandSegmentSizes"}); p << " : "; - for (auto val : getValues()) + for (auto val : getValues()) { p << val.getType() << ", "; + } p << getDestination().getType(); - if (!getMask().empty()) + if (!getMask().empty()) { p << ", " << getMask()[0].getType(); + } } LogicalResult VMIvStoreOp::verify() { // group and dist_mode are mutually exclusive - if (getGroup() && getDistMode()) + if (getGroup() && getDistMode()) { return emitOpError("group and dist_mode are mutually exclusive"); - if (getGroup() && !getStride()) + } + if (getGroup() && !getStride()) { return emitOpError("group requires a stride operand"); - if (!getGroup() && getStride()) + } + if (!getGroup() && getStride()) { return emitOpError("stride operand is only valid with group"); - if (getGroup() && !getMask().empty()) + } + if (getGroup() && !getMask().empty()) { return emitOpError("group mode does not support mask operand"); + } if (getGroup()) { int64_t numGroups = getGroupAttr().getInt(); - if (numGroups <= 0) + if (numGroups <= 0) { return emitOpError("group must be positive, got ") << numGroups; - if (getValues().size() != 1) + } + if (getValues().size() != 1) { return emitOpError("group mode requires exactly 1 value"); + } auto valueType = cast(getValues()[0].getType()); - if (failed(verifyNumGroups(getOperation(), valueType, numGroups))) + if (failed(verifyNumGroups(getOperation(), valueType, numGroups))) { return failure(); + } } // block_stride / repeat_stride: paired, mutually exclusive with @@ -4004,34 +4293,40 @@ LogicalResult VMIvStoreOp::verify() { if (getDistMode()) return emitOpError( "block_stride and dist_mode are mutually exclusive"); - if (getValues().size() != 1) + if (getValues().size() != 1) { return emitOpError("block-stride mode requires exactly 1 value"); + } } auto distMode = getDistMode(); bool isDintlv = distMode && *distMode == "dintlv"; size_t nValues = getValues().size(); - if (nValues < 1) + if (nValues < 1) { return emitOpError("requires at least 1 value"); - if (isDintlv && nValues != 2) + } + if (isDintlv && nValues != 2) { return emitOpError("dist-mode \"dintlv\" requires exactly 2 values"); + } if (!isDintlv && nValues != 1) return emitOpError("requires exactly 1 value for dist-mode \"") << (distMode ? *distMode : "continuous") << "\""; bool hasMask = !getMask().empty(); - if (getMask().size() > 1) + if (getMask().size() > 1) { return emitOpError("at most one mask allowed"); + } - if (distMode && !validDistModes().count(*distMode)) + if (distMode && !validDistModes().count(*distMode)) { return emitOpError("invalid dist-mode: \"") << *distMode << "\""; + } if (distMode && (*distMode == "unpack" || *distMode == "brc")) return emitOpError("dist-mode \"") << *distMode << "\" is not valid for vstore"; auto pmode = getPmode(); - if (pmode && !validPModes().count(*pmode)) + if (pmode && !validPModes().count(*pmode)) { return emitOpError("invalid pmode: \"") << *pmode << "\""; + } if (pmode && *pmode != "zero") return emitOpError("pmode \"merge\" is not supported for stores: the " "legacy store lowering is mask-governed only and " @@ -4064,8 +4359,9 @@ LogicalResult VMIvStoreOp::verify() { if (hasMask) { auto maskType = cast(getMask()[0].getType()); - if (failed(verifyMaskMatchesData(getOperation(), maskType, valueType))) + if (failed(verifyMaskMatchesData(getOperation(), maskType, valueType))) { return failure(); + } } return success(); @@ -4086,8 +4382,9 @@ LogicalResult VMIVsstbOp::verify() { failed(verifyUBBackedMemory(getOperation(), getDestination().getType(), "destination"))) return failure(); - if (auto pmode = getPmode(); pmode && !validPModes().count(*pmode)) + if (auto pmode = getPmode(); pmode && !validPModes().count(*pmode)) { return emitOpError("invalid pmode: \"") << *pmode << "\""; + } if (auto pmode = getPmode(); pmode && *pmode != "zero") return emitOpError("pmode \"merge\" is not supported for stores: the " "legacy store lowering is mask-governed only and " @@ -4115,8 +4412,9 @@ LogicalResult VMIVselrOp::verify() { return emitOpError( "requires index lane count to match result lane count"); - if (!isa(indexType.getElementType())) + if (!isa(indexType.getElementType())) { return emitOpError("requires index element type to be integer"); + } unsigned sourceBits = pto::getPTOStorageElemBitWidth(sourceType.getElementType()); @@ -4211,8 +4509,9 @@ LogicalResult VMIVcmpOp::verify() { } // Seed mask must match data shape. - if (failed(verifyMaskMatchesData(getOperation(), seedType, lhsType))) + if (failed(verifyMaskMatchesData(getOperation(), seedType, lhsType))) { return failure(); + } // Result mask must match seed mask. if (seedType.getElementCount() != resultType.getElementCount()) @@ -4257,8 +4556,9 @@ LogicalResult VMIVcmpsOp::verify() { } // Seed mask must match data shape. - if (failed(verifyMaskMatchesData(getOperation(), seedType, srcType))) + if (failed(verifyMaskMatchesData(getOperation(), seedType, srcType))) { return failure(); + } // Result mask must match seed mask. if (seedType.getElementCount() != resultType.getElementCount()) @@ -4290,29 +4590,35 @@ ParseResult VMIvLoadOp::parse(OpAsmParser &parser, OperationState &result) { int numPostBracket = 0; OpAsmParser::UnresolvedOperand postOp1, postOp2; if (succeeded(parser.parseOptionalComma())) { - if (parser.parseOperand(postOp1)) + if (parser.parseOperand(postOp1)) { return failure(); + } numPostBracket = 1; if (succeeded(parser.parseOptionalComma())) { - if (parser.parseOperand(postOp2)) + if (parser.parseOperand(postOp2)) { return failure(); + } numPostBracket = 2; } } - if (parser.parseOptionalAttrDict(result.attributes)) + if (parser.parseOptionalAttrDict(result.attributes)) { return failure(); + } Type sourceType; - if (parser.parseColonType(sourceType)) + if (parser.parseColonType(sourceType)) { return failure(); + } - if (parser.parseArrow()) + if (parser.parseArrow()) { return failure(); + } SmallVector resultTypes; - if (parser.parseTypeList(resultTypes)) + if (parser.parseTypeList(resultTypes)) { return failure(); + } // Disambiguate post-bracket operands bool hasStride = false; @@ -4333,8 +4639,9 @@ ParseResult VMIvLoadOp::parse(OpAsmParser &parser, OperationState &result) { strideOperand = postOp1; } - if (parser.resolveOperand(sourceOperand, sourceType, result.operands)) + if (parser.resolveOperand(sourceOperand, sourceType, result.operands)) { return failure(); + } if (parser.resolveOperand(offsetOperand, parser.getBuilder().getIndexType(), result.operands)) return failure(); @@ -4383,22 +4690,28 @@ void VMIvLoadOp::print(OpAsmPrinter &p) { LogicalResult VMIvLoadOp::verify() { // group and dist_mode are mutually exclusive, except brc which supports // group broadcast (one scalar per group → broadcast within each group). - if (getGroup() && getDistMode() && getDistMode() != "brc") + if (getGroup() && getDistMode() && getDistMode() != "brc") { return emitOpError("group and dist_mode are mutually exclusive"); - if (getGroup() && !getStride()) + } + if (getGroup() && !getStride()) { return emitOpError("group requires a stride operand"); - if (!getGroup() && getStride()) + } + if (!getGroup() && getStride()) { return emitOpError("stride operand is only valid with group"); + } if (getGroup()) { int64_t numGroups = getGroupAttr().getInt(); - if (numGroups <= 0) + if (numGroups <= 0) { return emitOpError("group must be positive, got ") << numGroups; - if (getResults().size() != 1) + } + if (getResults().size() != 1) { return emitOpError("group mode requires exactly 1 result"); + } auto resultType = cast(getResults()[0].getType()); - if (failed(verifyNumGroups(getOperation(), resultType, numGroups))) + if (failed(verifyNumGroups(getOperation(), resultType, numGroups))) { return failure(); + } } // block_stride and repeat_stride must be paired, mutually exclusive @@ -4412,25 +4725,29 @@ LogicalResult VMIvLoadOp::verify() { if (getDistMode()) return emitOpError( "block_stride and dist_mode are mutually exclusive"); - if (getResults().size() != 1) + if (getResults().size() != 1) { return emitOpError("block-stride mode requires exactly 1 result"); + } } // result count vs dist-mode auto distMode = getDistMode(); bool isDintlv = distMode && *distMode == "dintlv"; size_t nResults = getResults().size(); - if (isDintlv && nResults != 2) + if (isDintlv && nResults != 2) { return emitOpError("dist-mode \"dintlv\" requires exactly 2 results"); + } if (!isDintlv && nResults != 1) return emitOpError("requires exactly 1 result for dist-mode \"") << (distMode ? *distMode : "continuous") << "\""; - if (distMode && !validDistModes().count(*distMode)) + if (distMode && !validDistModes().count(*distMode)) { return emitOpError("invalid dist-mode: \"") << *distMode << "\""; + } auto pmode = getPmode(); - if (pmode && !validPModes().count(*pmode)) + if (pmode && !validPModes().count(*pmode)) { return emitOpError("invalid pmode: \"") << *pmode << "\""; + } bool isUnpack = distMode && *distMode == "unpack"; for (auto res : getResults()) { @@ -4469,30 +4786,37 @@ Type mlir::pto::getVMIPhysicalDataElementType(VMIVRegType type) { FailureOr mlir::pto::getDataLanesPerPart(Type elementType) { unsigned elementBitWidth = pto::getPTOStorageElemBitWidth(elementType); - if (elementBitWidth == 0) + if (elementBitWidth == 0) { return failure(); + } constexpr int64_t kPhysicalVRegBits = 256 * 8; - if (kPhysicalVRegBits % elementBitWidth != 0) + if (kPhysicalVRegBits % elementBitWidth != 0) { return failure(); + } return kPhysicalVRegBits / elementBitWidth; } FailureOr mlir::pto::getMaskLanesPerPart(StringRef granularity) { - if (granularity == "b8") + if (granularity == "b8") { return 256; - if (granularity == "b16") + } + if (granularity == "b16") { return 128; - if (granularity == "b32") + } + if (granularity == "b32") { return 64; + } return failure(); } FailureOr mlir::pto::getVMILayoutBlockElems(Type type) { FailureOr layout = getAssignedVMILayout(type); - if (failed(layout)) + if (failed(layout)) { return failure(); - if (!(*layout).isBlockDeinterleaved()) + } + if (!(*layout).isBlockDeinterleaved()) { return 1; + } FailureOr lanesPerPart = getPhysicalLanesPerPart(type); constexpr int64_t kVCGBlocksPerPart = 8; @@ -4506,8 +4830,9 @@ FailureOr mlir::pto::getVMIPhysicalArity(Type type) { FailureOr elementCount = getVMIElementCount(type); FailureOr lanesPerPart = getPhysicalLanesPerPart(type); FailureOr layout = getAssignedVMILayout(type); - if (failed(elementCount) || failed(lanesPerPart) || failed(layout)) + if (failed(elementCount) || failed(lanesPerPart) || failed(layout)) { return failure(); + } if ((*layout).isGroupSlots() && (*layout).getSlots() > 0) return divideCeilNonNegative((*layout).getNumGroups(), @@ -4515,8 +4840,9 @@ FailureOr mlir::pto::getVMIPhysicalArity(Type type) { int64_t factor = (*layout).isDenseSplit() ? (*layout).getFactor() : 1; FailureOr blockElems = getVMILayoutBlockElems(type); - if (failed(blockElems)) + if (failed(blockElems)) { return failure(); + } int64_t laneStride = isa(type) ? 1 : ((*layout).isDense() ? (*layout).getLaneStride() @@ -4542,24 +4868,27 @@ mlir::pto::mapLogicalLaneToPhysical(Type type, int64_t logicalLane) { if (failed(elementCount) || failed(factor) || failed(blockElems) || failed(laneStride) || failed(lanesPerPart)) return failure(); - if (logicalLane < 0 || logicalLane >= *elementCount) + if (logicalLane < 0 || logicalLane >= *elementCount) { return failure(); + } FailureOr layout = getAssignedVMILayout(type); if (succeeded(layout) && (*layout).isGroupSlots() && (*layout).getSlots() > 0) { int64_t slots = (*layout).getSlots(); int64_t lane = logicalLane % slots; - if (lane >= *lanesPerPart) + if (lane >= *lanesPerPart) { return failure(); + } return VMIPhysicalLane{/*part=*/0, logicalLane / slots, lane}; } int64_t part = 0; std::optional indexInPart = mapDenseLogicalLaneToPartIndex( *elementCount, *factor, *blockElems, logicalLane, part); - if (!indexInPart) + if (!indexInPart) { return failure(); + } int64_t physicalIndex = *indexInPart * *laneStride; return VMIPhysicalLane{part, physicalIndex / *lanesPerPart, physicalIndex % *lanesPerPart}; @@ -4584,22 +4913,26 @@ FailureOr mlir::pto::mapPhysicalLaneToLogical(Type type, int64_t part, if (succeeded(layout) && (*layout).isGroupSlots() && (*layout).getSlots() > 0) { int64_t slots = (*layout).getSlots(); - if (part != 0 || lane >= slots) + if (part != 0 || lane >= slots) { return failure(); + } int64_t logicalLane = chunk * slots + lane; - if (logicalLane >= *elementCount) + if (logicalLane >= *elementCount) { return failure(); + } return logicalLane; } int64_t physicalIndexInPart = chunk * *lanesPerPart + lane; - if (physicalIndexInPart % *laneStride != 0) + if (physicalIndexInPart % *laneStride != 0) { return failure(); + } int64_t indexInPart = physicalIndexInPart / *laneStride; std::optional logicalLane = mapDensePartIndexToLogicalLane( *elementCount, *factor, *blockElems, part, indexInPart); - if (!logicalLane) + if (!logicalLane) { return failure(); + } return *logicalLane; } @@ -4621,18 +4954,21 @@ FailureOr mlir::pto::isPaddingLane(Type type, int64_t part, int64_t chunk, if (succeeded(layout) && (*layout).isGroupSlots() && (*layout).getSlots() > 0) { int64_t slots = (*layout).getSlots(); - if (part != 0) + if (part != 0) { return true; - if (lane >= slots) + } + if (lane >= slots) { return true; + } return chunk * slots + lane >= *elementCount; } int64_t lanesInPart = getDenseLogicalLanesInPart(*elementCount, *factor, *blockElems, part); int64_t physicalIndexInPart = chunk * *lanesPerPart + lane; - if (physicalIndexInPart % *laneStride != 0) + if (physicalIndexInPart % *laneStride != 0) { return true; + } int64_t indexInPart = physicalIndexInPart / *laneStride; return indexInPart >= lanesInPart; } diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index f7352264ba..a63c4bc468 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -64,8 +64,9 @@ static std::string formatMaskType(StringRef granularity) { static LogicalResult verifyVRegTypeLike(Operation *op, Type type, StringRef roleDescription) { auto vecType = dyn_cast(type); - if (!vecType) + if (!vecType) { return op->emitOpError() << roleDescription << " must be !pto.vreg<...>"; + } return VRegType::verify( [&]() { return op->emitOpError() << roleDescription << " "; }, @@ -74,16 +75,18 @@ static LogicalResult verifyVRegTypeLike(Operation *op, Type type, static LogicalResult verifyMaskTypeLike(Operation *op, Type type, StringRef roleDescription) { - if (!isa(type)) + if (!isa(type)) { return op->emitOpError() << roleDescription << " must be !pto.mask<...>"; + } return success(); } static LogicalResult verifyNonLowPrecisionVRegElementTypeLike( Operation *op, Type type, StringRef roleDescription) { auto vecType = dyn_cast(type); - if (!vecType) + if (!vecType) { return success(); + } if (pto::isPTOLowPrecisionType(vecType.getElementType())) return op->emitOpError() << roleDescription @@ -97,8 +100,9 @@ static LogicalResult verifyMaskTypeWithGranularityLike(Operation *op, Type type, StringRef roleDescription, StringRef granularity) { auto maskType = dyn_cast(type); - if (!maskType) + if (!maskType) { return op->emitOpError() << roleDescription << " must be !pto.mask<...>"; + } if (maskType.getGranularity() != granularity) { return op->emitOpError() << roleDescription << " must be " << formatMaskType(granularity); @@ -139,8 +143,9 @@ static bool isMaskGranularityAdjacentNarrowing(StringRef inputGranularity, } static bool isSupportedShuffleValueType(Type type) { - if (auto intType = dyn_cast(type)) + if (auto intType = dyn_cast(type)) { return intType.getWidth() == 32 || intType.getWidth() == 64; + } if (auto vecType = dyn_cast(type)) return vecType.getRank() == 1 && vecType.getDimSize(0) == 2 && vecType.getElementType().isF16(); @@ -148,8 +153,9 @@ static bool isSupportedShuffleValueType(Type type) { } static bool isSupportedReduxValueType(Type type) { - if (auto intType = dyn_cast(type)) + if (auto intType = dyn_cast(type)) { return intType.getWidth() == 32; + } return type.isF16() || type.isF32(); } @@ -176,11 +182,13 @@ LogicalResult SimtLaunchOp::verify() { } FunctionType calleeType = callee.getFunctionType(); - if (!calleeType.getResults().empty()) + if (!calleeType.getResults().empty()) { return emitOpError("requires a callee with no results"); + } - if (calleeType.getNumInputs() != getArgs().size()) + if (calleeType.getNumInputs() != getArgs().size()) { return emitOpError("incorrect number of operands for callee"); + } for (auto [index, argType, operand] : llvm::enumerate(calleeType.getInputs(), getArgs())) { @@ -205,8 +213,9 @@ static LogicalResult verifyShuffleSemanticControl(Operation *op, << " operand to be i32"; int64_t width = widthAttr.getInt(); - if (width != 16 && width != 32) + if (width != 16 && width != 32) { return op->emitOpError() << "requires width to be 16 or 32"; + } return success(); } @@ -229,8 +238,9 @@ static LogicalResult verifyReduxSemanticType(Operation *op, Type valueType, return op->emitOpError() << "requires explicit signedness for integer redux"; - if (!signednessAttr) + if (!signednessAttr) { return success(); + } auto signedness = cast(signednessAttr).getValue(); (void)signedness; @@ -238,8 +248,9 @@ static LogicalResult verifyReduxSemanticType(Operation *op, Type valueType, } static bool isStandardScalarConvertType(Type type) { - if (auto intType = dyn_cast(type)) + if (auto intType = dyn_cast(type)) { return intType.getWidth() == 32 || intType.getWidth() == 64; + } return type.isF16() || type.isBF16() || type.isF32(); } @@ -254,8 +265,9 @@ static bool isVector2Of(Type type, llvm::function_ref elementPred) { } static bool isSupportedPackedConvertType(Type type) { - if (pto::isPTOHiFloat8x2Type(type)) + if (pto::isPTOHiFloat8x2Type(type)) { return true; + } return isVector2Of(type, [](Type elem) { return elem.isF16() || elem.isBF16() || elem.isF32() || pto::isPTOFloat8Type(elem) || pto::isPTOHiFloat8Type(elem); @@ -464,8 +476,9 @@ static LogicalResult verifyConvertControls(Operation *op, Type srcType, } static bool isSupportedAtomicScalarType(Type type) { - if (auto intType = dyn_cast(type)) + if (auto intType = dyn_cast(type)) { return intType.getWidth() == 32 || intType.getWidth() == 64; + } return type.isF16() || type.isBF16() || type.isF32() || isVector2F16OrBF16Type(type); } @@ -482,24 +495,29 @@ static LogicalResult verifyAtomicCommon(Operation *op, Value ptr, Type valueType << "requires atomic result type to match value type"; auto ptrTy = dyn_cast(ptr.getType()); - if (!ptrTy) + if (!ptrTy) { return op->emitOpError() << "requires !pto.ptr pointer operand"; + } if (ptrTy.getElementType() != valueType) return op->emitOpError() << "requires atomic value type to match pointer element type"; AddressSpace addressSpace = ptrTy.getMemorySpace().getAddressSpace(); - if (addressSpace != AddressSpace::GM && addressSpace != AddressSpace::VEC) + if (addressSpace != AddressSpace::GM && addressSpace != AddressSpace::VEC) { return op->emitOpError() << "requires GM or UB pointer"; - if (addressSpace == AddressSpace::VEC && valueType.isInteger(64)) + } + if (addressSpace == AddressSpace::VEC && valueType.isInteger(64)) { return op->emitOpError() << "does not support i64 UB-space atomics"; + } auto intType = dyn_cast(valueType); if (bitwise) { - if (!intType) + if (!intType) { return op->emitOpError() << "requires integer type for bitwise atomics"; - if (addressSpace == AddressSpace::VEC && intType.getWidth() == 64) + } + if (addressSpace == AddressSpace::VEC && intType.getWidth() == 64) { return op->emitOpError() << "does not support i64 UB-space bitwise atomics"; + } } if (signednessAttr && !intType) @@ -521,23 +539,28 @@ static LogicalResult verifyAtomicCommon(Operation *op, Value ptr, Type valueType static LogicalResult verifyLdgStgAccess(Operation *op, Type ptrType, Type valueType) { auto ptrTy = dyn_cast(ptrType); - if (!ptrTy) + if (!ptrTy) { return op->emitOpError() << "requires !pto.ptr operand"; - if (ptrTy.getMemorySpace().getAddressSpace() != AddressSpace::GM) + } + if (ptrTy.getMemorySpace().getAddressSpace() != AddressSpace::GM) { return op->emitOpError() << "requires GM pointer"; + } if (auto intType = dyn_cast(valueType)) { unsigned width = intType.getWidth(); - if (width == 8 || width == 16 || width == 32 || width == 64) + if (width == 8 || width == 16 || width == 32 || width == 64) { return success(); + } } if (valueType.isF16() || valueType.isBF16() || valueType.isF32() || valueType.isF64()) return success(); - if (pto::isPTOFloat8Type(valueType) || pto::isPTOHiFloat8Type(valueType)) + if (pto::isPTOFloat8Type(valueType) || pto::isPTOHiFloat8Type(valueType)) { return success(); - if (pto::isPTOPackedLdgStgVectorType(valueType)) + } + if (pto::isPTOPackedLdgStgVectorType(valueType)) { return success(); + } return op->emitOpError() << "currently supports 8/16/32/64-bit integer, " @@ -553,10 +576,12 @@ static LogicalResult verifyLdStDevAccess(Operation *op, Type ptrType, << "does not accept l1cache or l2cache policy attributes"; auto ptrTy = dyn_cast(ptrType); - if (!ptrTy) + if (!ptrTy) { return op->emitOpError() << "requires !pto.ptr operand"; - if (ptrTy.getMemorySpace().getAddressSpace() != AddressSpace::GM) + } + if (ptrTy.getMemorySpace().getAddressSpace() != AddressSpace::GM) { return op->emitOpError() << "requires GM pointer"; + } auto intType = dyn_cast(valueType); if (!intType || (intType.getWidth() != 8 && intType.getWidth() != 16 && @@ -674,8 +699,9 @@ LogicalResult MulhiOp::verify() { LogicalResult MulI32ToI64Op::verify() { return success(); } LogicalResult AtomicCasOp::verify() { - if (getCompare().getType() != getValue().getType()) + if (getCompare().getType() != getValue().getType()) { return emitOpError() << "requires compare and value types to match"; + } return verifyAtomicCommon(getOperation(), getPtr(), getValue().getType(), getOld().getType(), /*bitwise=*/false, getSignednessAttr()); @@ -839,8 +865,9 @@ static LogicalResult verifyNotNestedInVecScope(Operation *op, static LogicalResult verifyNestedInVecScope(Operation *op, StringRef opNameForDiag) { - if (op->getParentOfType() || op->getParentOfType()) + if (op->getParentOfType() || op->getParentOfType()) { return success(); + } return op->emitOpError() << "must be nested under pto.vecscope/pto.strict_vecscope; " << opNameForDiag << " is part of the vecscope control sequence"; @@ -848,8 +875,9 @@ static LogicalResult verifyNestedInVecScope(Operation *op, static LogicalResult verifyAlignTypeLike(Operation *op, Type type, StringRef roleDescription) { - if (!isa(type)) + if (!isa(type)) { return op->emitOpError() << roleDescription << " must be !pto.align"; + } return success(); } @@ -862,8 +890,9 @@ static bool isSupportedMovPadScalarType(Type type) { return intType.isSignless() && (intType.getWidth() == 8 || intType.getWidth() == 16 || intType.getWidth() == 32); - if (auto floatType = dyn_cast(type)) + if (auto floatType = dyn_cast(type)) { return floatType.isF16() || floatType.isBF16() || floatType.isF32(); + } return false; } @@ -885,10 +914,12 @@ static std::optional getVdupMaskGranularity(Type elementType) { return std::nullopt; } } - if (elementType.isF16() || elementType.isBF16()) + if (elementType.isF16() || elementType.isBF16()) { return StringRef("b16"); - if (elementType.isF32()) + } + if (elementType.isF32()) { return StringRef("b32"); + } return std::nullopt; } @@ -912,20 +943,24 @@ static bool isLoadAlignProducer(Operation *op) { static scf::IfOp getEnclosingBranchIf(Operation *op) { for (Operation *cursor = op; cursor; cursor = cursor->getParentOp()) { auto ifOp = dyn_cast(cursor); - if (!ifOp) + if (!ifOp) { continue; + } Region *parentRegion = op->getParentRegion(); - if (parentRegion == &ifOp.getThenRegion() || parentRegion == &ifOp.getElseRegion()) + if (parentRegion == &ifOp.getThenRegion() || parentRegion == &ifOp.getElseRegion()) { return ifOp; + } } return nullptr; } static bool isValueOwnedByRegion(Value value, Region *region) { - if (auto blockArg = dyn_cast(value)) + if (auto blockArg = dyn_cast(value)) { return blockArg.getParentRegion() == region; - if (Operation *def = value.getDefiningOp()) + } + if (Operation *def = value.getDefiningOp()) { return def->getParentRegion() == region; + } return false; } @@ -934,7 +969,6 @@ static FailureOr resolveLoadAlignRoot(Value value, Operation *user); static FailureOr resolveStoreAlignRootImpl( Value current, llvm::SmallPtrSet visited) { - while (true) { if (!visited.insert(current.getAsOpaquePointer()).second) { return failure(); @@ -943,22 +977,26 @@ static FailureOr resolveStoreAlignRootImpl( if (auto blockArg = dyn_cast(current)) { auto *owner = blockArg.getOwner(); auto forOp = dyn_cast(owner->getParentOp()); - if (!forOp) + if (!forOp) { return failure(); + } unsigned argNumber = blockArg.getArgNumber(); unsigned ivCount = forOp.getNumInductionVars(); - if (argNumber < ivCount) + if (argNumber < ivCount) { return failure(); + } unsigned iterIdx = argNumber - ivCount; - if (iterIdx >= forOp.getInitArgs().size()) + if (iterIdx >= forOp.getInitArgs().size()) { return failure(); + } current = forOp.getInitArgs()[iterIdx]; continue; } if (Operation *def = current.getDefiningOp()) { - if (isa(def)) + if (isa(def)) { return current; + } if (auto stateOp = dyn_cast(def)) { current = stateOp.getAlignIn(); continue; @@ -973,18 +1011,21 @@ static FailureOr resolveStoreAlignRootImpl( } if (auto forOp = dyn_cast(def)) { auto result = dyn_cast(current); - if (!result) + if (!result) { return failure(); + } unsigned resultIdx = result.getResultNumber(); - if (resultIdx >= forOp.getYieldedValues().size()) + if (resultIdx >= forOp.getYieldedValues().size()) { return failure(); + } current = forOp.getYieldedValues()[resultIdx]; continue; } if (auto ifOp = dyn_cast(def)) { auto result = dyn_cast(current); - if (!result || !ifOp.elseBlock()) + if (!result || !ifOp.elseBlock()) { return failure(); + } unsigned resultIdx = result.getResultNumber(); auto thenYield = dyn_cast(ifOp.thenBlock()->getTerminator()); auto elseYield = dyn_cast(ifOp.elseBlock()->getTerminator()); @@ -996,8 +1037,9 @@ static FailureOr resolveStoreAlignRootImpl( resolveStoreAlignRootImpl(thenYield.getOperand(resultIdx), visited); FailureOr elseRoot = resolveStoreAlignRootImpl(elseYield.getOperand(resultIdx), visited); - if (failed(thenRoot) || failed(elseRoot) || *thenRoot != *elseRoot) + if (failed(thenRoot) || failed(elseRoot) || *thenRoot != *elseRoot) { return failure(); + } return *thenRoot; } } @@ -1016,8 +1058,9 @@ static LogicalResult verifyStoreAlignLoopThreading(Value align, Operation *user, Operation *cursor = user; while (auto forOp = cursor->getParentOfType()) { Region *body = &forOp.getRegion(); - if (isValueOwnedByRegion(align, body)) + if (isValueOwnedByRegion(align, body)) { return success(); + } if (!isValueOwnedByRegion(align, body)) { return user->emitOpError() << roleDescription @@ -1032,11 +1075,13 @@ static LogicalResult verifyStoreAlignLoopThreading(Value align, Operation *user, static FailureOr resolveSingleAlignIfResult(scf::IfOp ifOp) { SmallVector alignResultIndices; for (auto [index, type] : llvm::enumerate(ifOp.getResultTypes())) { - if (isa(type)) + if (isa(type)) { alignResultIndices.push_back(index); + } } - if (alignResultIndices.size() != 1) + if (alignResultIndices.size() != 1) { return failure(); + } return ifOp.getResult(alignResultIndices.front()); } @@ -1107,8 +1152,9 @@ static LogicalResult verifyStoreAlignLinearUses(Value value, Operation *user) { commonIf = nullptr; break; } - if (!commonIf) + if (!commonIf) { commonIf = enclosingIf; + } else if (commonIf != enclosingIf) { commonIf = nullptr; break; @@ -1124,8 +1170,9 @@ static LogicalResult verifyStoreAlignLinearUses(Value value, Operation *user) { return user->emitOpError() << "!pto.align value must form a single linear store-state chain"; } - if (nextValues.empty()) + if (nextValues.empty()) { return success(); + } current = nextValues.front(); } @@ -1134,14 +1181,17 @@ static LogicalResult verifyStoreAlignLinearUses(Value value, Operation *user) { static LogicalResult verifyStoreAlignChain(Value align, Operation *user, StringRef roleDescription) { - if (disableVPTOAlignChainVerification) + if (disableVPTOAlignChainVerification) { return success(); + } - if (failed(verifyAlignTypeLike(user, align.getType(), roleDescription))) + if (failed(verifyAlignTypeLike(user, align.getType(), roleDescription))) { return failure(); + } - if (failed(verifyStoreAlignLoopThreading(align, user, roleDescription))) + if (failed(verifyStoreAlignLoopThreading(align, user, roleDescription))) { return failure(); + } FailureOr root = resolveStoreAlignRoot(align, user); if (failed(root)) { @@ -1171,48 +1221,55 @@ static LogicalResult verifyStoreAlignChain(Value align, Operation *user, static FailureOr resolveLoadAlignRootImpl( Value current, llvm::SmallPtrSet visited) { - while (true) { - if (!visited.insert(current.getAsOpaquePointer()).second) + if (!visited.insert(current.getAsOpaquePointer()).second) { return failure(); + } if (auto blockArg = dyn_cast(current)) { auto *owner = blockArg.getOwner(); auto forOp = dyn_cast(owner->getParentOp()); - if (!forOp) + if (!forOp) { return failure(); + } unsigned argNumber = blockArg.getArgNumber(); unsigned ivCount = forOp.getNumInductionVars(); - if (argNumber < ivCount) + if (argNumber < ivCount) { return failure(); + } unsigned iterIdx = argNumber - ivCount; - if (iterIdx >= forOp.getInitArgs().size()) + if (iterIdx >= forOp.getInitArgs().size()) { return failure(); + } current = forOp.getInitArgs()[iterIdx]; continue; } if (Operation *def = current.getDefiningOp()) { - if (isa(def)) + if (isa(def)) { return current; + } if (auto stateOp = dyn_cast(def)) { current = stateOp.getAlign(); continue; } if (auto forOp = dyn_cast(def)) { auto result = dyn_cast(current); - if (!result) + if (!result) { return failure(); + } unsigned resultIdx = result.getResultNumber(); - if (resultIdx >= forOp.getYieldedValues().size()) + if (resultIdx >= forOp.getYieldedValues().size()) { return failure(); + } current = forOp.getYieldedValues()[resultIdx]; continue; } if (auto ifOp = dyn_cast(def)) { auto result = dyn_cast(current); - if (!result || !ifOp.elseBlock()) + if (!result || !ifOp.elseBlock()) { return failure(); + } unsigned resultIdx = result.getResultNumber(); auto thenYield = dyn_cast(ifOp.thenBlock()->getTerminator()); auto elseYield = dyn_cast(ifOp.elseBlock()->getTerminator()); @@ -1224,8 +1281,9 @@ static FailureOr resolveLoadAlignRootImpl( resolveLoadAlignRootImpl(thenYield.getOperand(resultIdx), visited); FailureOr elseRoot = resolveLoadAlignRootImpl(elseYield.getOperand(resultIdx), visited); - if (failed(thenRoot) || failed(elseRoot) || *thenRoot != *elseRoot) + if (failed(thenRoot) || failed(elseRoot) || *thenRoot != *elseRoot) { return failure(); + } return *thenRoot; } } @@ -1244,8 +1302,9 @@ static LogicalResult verifyLoadAlignLoopThreading(Value align, Operation *user, Operation *cursor = user; while (auto forOp = cursor->getParentOfType()) { Region *body = &forOp.getRegion(); - if (isValueOwnedByRegion(align, body)) + if (isValueOwnedByRegion(align, body)) { return success(); + } if (!isValueOwnedByRegion(align, body)) { return user->emitOpError() << roleDescription @@ -1312,8 +1371,9 @@ static LogicalResult verifyLoadAlignLinearUses(Value value, Operation *user) { commonIf = nullptr; break; } - if (!commonIf) + if (!commonIf) { commonIf = enclosingIf; + } else if (commonIf != enclosingIf) { commonIf = nullptr; break; @@ -1329,8 +1389,9 @@ static LogicalResult verifyLoadAlignLinearUses(Value value, Operation *user) { return user->emitOpError() << "!pto.align value must form a single linear load-state chain"; } - if (nextValues.empty()) + if (nextValues.empty()) { return success(); + } current = nextValues.front(); } @@ -1339,14 +1400,17 @@ static LogicalResult verifyLoadAlignLinearUses(Value value, Operation *user) { static LogicalResult verifyLoadAlignChain(Value align, Operation *user, StringRef roleDescription) { - if (disableVPTOAlignChainVerification) + if (disableVPTOAlignChainVerification) { return success(); + } - if (failed(verifyAlignTypeLike(user, align.getType(), roleDescription))) + if (failed(verifyAlignTypeLike(user, align.getType(), roleDescription))) { return failure(); + } - if (failed(verifyLoadAlignLoopThreading(align, user, roleDescription))) + if (failed(verifyLoadAlignLoopThreading(align, user, roleDescription))) { return failure(); + } FailureOr root = resolveLoadAlignRoot(align, user); if (failed(root)) { @@ -1398,54 +1462,70 @@ static bool isSupportedPartToken(StringRef part) { static bool isSupportedSprToken(StringRef spr) { return spr == "AR"; } static std::optional normalizeRoundModeToken(StringRef token) { - if (token == "R" || token == "ROUND_R") + if (token == "R" || token == "ROUND_R") { return StringRef("R"); - if (token == "A" || token == "ROUND_A") + } + if (token == "A" || token == "ROUND_A") { return StringRef("A"); - if (token == "F" || token == "ROUND_F") + } + if (token == "F" || token == "ROUND_F") { return StringRef("F"); - if (token == "C" || token == "ROUND_C") + } + if (token == "C" || token == "ROUND_C") { return StringRef("C"); - if (token == "Z" || token == "ROUND_Z") + } + if (token == "Z" || token == "ROUND_Z") { return StringRef("Z"); - if (token == "O" || token == "ROUND_O") + } + if (token == "O" || token == "ROUND_O") { return StringRef("O"); - if (token == "H" || token == "ROUND_H") + } + if (token == "H" || token == "ROUND_H") { return StringRef("H"); + } return std::nullopt; } static std::optional normalizeSaturationToken(StringRef token) { - if (token == "SAT" || token == "RS_ENABLE") + if (token == "SAT" || token == "RS_ENABLE") { return StringRef("SAT"); - if (token == "NOSAT" || token == "RS_DISABLE") + } + if (token == "NOSAT" || token == "RS_DISABLE") { return StringRef("NOSAT"); + } return std::nullopt; } static std::optional normalizeEvenOddPartToken(StringRef token) { - if (token == "EVEN" || token == "PART_EVEN") + if (token == "EVEN" || token == "PART_EVEN") { return StringRef("EVEN"); - if (token == "ODD" || token == "PART_ODD") + } + if (token == "ODD" || token == "PART_ODD") { return StringRef("ODD"); + } return std::nullopt; } static std::optional normalizePacked4PartToken(StringRef token) { - if (token == "P0" || token == "PART_P0") + if (token == "P0" || token == "PART_P0") { return StringRef("P0"); - if (token == "P1" || token == "PART_P1") + } + if (token == "P1" || token == "PART_P1") { return StringRef("P1"); - if (token == "P2" || token == "PART_P2") + } + if (token == "P2" || token == "PART_P2") { return StringRef("P2"); - if (token == "P3" || token == "PART_P3") + } + if (token == "P3" || token == "PART_P3") { return StringRef("P3"); + } return std::nullopt; } static std::optional normalizeVcvtPartToken(StringRef token) { - if (auto normalized = normalizeEvenOddPartToken(token)) + if (auto normalized = normalizeEvenOddPartToken(token)) { return normalized; + } return normalizePacked4PartToken(token); } @@ -1484,22 +1564,30 @@ struct VcvtContract { }; static VcvtElemKind classifyVcvtElemType(Type type) { - if (type.isF16()) + if (type.isF16()) { return VcvtElemKind::F16; - if (type.isBF16()) + } + if (type.isBF16()) { return VcvtElemKind::BF16; - if (type.isF32()) + } + if (type.isF32()) { return VcvtElemKind::F32; - if (pto::isPTOFloat8E4M3LikeType(type)) + } + if (pto::isPTOFloat8E4M3LikeType(type)) { return VcvtElemKind::F8E4M3; - if (pto::isPTOFloat8E5M2LikeType(type)) + } + if (pto::isPTOFloat8E5M2LikeType(type)) { return VcvtElemKind::F8E5M2; - if (pto::isPTOHiFloat8Type(type)) + } + if (pto::isPTOHiFloat8Type(type)) { return VcvtElemKind::HiF8; - if (isa(type)) + } + if (isa(type)) { return VcvtElemKind::F4E1M2x2; - if (isa(type)) + } + if (isa(type)) { return VcvtElemKind::F4E2M1x2; + } if (auto intType = dyn_cast(type)) { switch (intType.getWidth()) { case 8: @@ -1548,10 +1636,12 @@ static std::optional classifyVcvtPartFamily(unsigned srcBits, unsigned dstBits) { unsigned largerBits = std::max(srcBits, dstBits); unsigned smallerBits = std::min(srcBits, dstBits); - if (largerBits == smallerBits * 2) + if (largerBits == smallerBits * 2) { return VcvtPartFamily::EvenOdd; - if (largerBits == smallerBits * 4) + } + if (largerBits == smallerBits * 4) { return VcvtPartFamily::Packed4; + } return std::nullopt; } @@ -1567,8 +1657,9 @@ static bool isValidVcvtPartForFamily(StringRef part, VcvtPartFamily family) { static bool isValidVcvtRoundModeForContract(StringRef roundMode, const VcvtContract &contract) { - if (!contract.allowedRndModes) + if (!contract.allowedRndModes) { return true; + } return StringRef(contract.allowedRndModes).contains(roundMode); } @@ -1758,14 +1849,18 @@ static std::optional lookupVcvtContract(VcvtElemKind src, } // namespace static std::optional getDistElementWidth(Type type) { - if (auto intType = dyn_cast(type)) + if (auto intType = dyn_cast(type)) { return intType.getWidth(); - if (type.isF16() || type.isBF16()) + } + if (type.isF16() || type.isBF16()) { return 16; - if (type.isF32()) + } + if (type.isF32()) { return 32; - if (type.isF64()) + } + if (type.isF64()) { return 64; + } return std::nullopt; } @@ -1798,23 +1893,31 @@ static bool isSupportedVstsx2DistToken(StringRef dist) { static std::optional getVstsMaskGranularityOverride(StringRef dist, Type elementType) { auto width = getDistElementWidth(elementType); - if (!width) + if (!width) { return std::nullopt; + } - if (dist == "MRG4CHN_B8") + if (dist == "MRG4CHN_B8") { return StringRef("b32"); - if (dist == "MRG2CHN_B8") + } + if (dist == "MRG2CHN_B8") { return StringRef("b16"); - if (dist == "MRG2CHN_B16") + } + if (dist == "MRG2CHN_B16") { return StringRef("b32"); - if (dist == "PK_B16") + } + if (dist == "PK_B16") { return StringRef("b16"); - if (dist == "PK_B32" || dist == "PK_B64" || dist == "PK4_B32") + } + if (dist == "PK_B32" || dist == "PK_B64" || dist == "PK4_B32") { return StringRef("b32"); - if (dist == "PK_B64") + } + if (dist == "PK_B64") { return StringRef("b32"); - if (dist == "PK4_B32") + } + if (dist == "PK4_B32") { return StringRef("b32"); + } return std::nullopt; } @@ -1824,15 +1927,18 @@ static bool isSupportedPostMode(StringRef mode) { } static unsigned getIntOrFloatBitWidth(Type type) { - if (auto intType = dyn_cast(type)) + if (auto intType = dyn_cast(type)) { return intType.getWidth(); - if (auto floatType = dyn_cast(type)) + } + if (auto floatType = dyn_cast(type)) { return floatType.getWidth(); + } if (pto::isPTOFloat8Type(type) || pto::isPTOHiFloat8Type(type) || pto::isPTOFloat4PackedType(type)) return 8; - if (pto::isPTOHiFloat8x2Type(type)) + if (pto::isPTOHiFloat8x2Type(type)) { return 16; + } return 0; } @@ -1842,18 +1948,21 @@ static bool isIntegerOrFloatLike(Type type) { static std::optional getVRegStorageBitWidth(Type type) { auto vecType = dyn_cast(type); - if (!vecType) + if (!vecType) { return std::nullopt; + } unsigned elemWidth = getIntOrFloatBitWidth(vecType.getElementType()); - if (!elemWidth) + if (!elemWidth) { return std::nullopt; + } return vecType.getElementCount() * static_cast(elemWidth); } static LogicalResult verifyIntegerVRegTypeLike(Operation *op, Type type, StringRef roleDescription) { - if (failed(verifyVRegTypeLike(op, type, roleDescription))) + if (failed(verifyVRegTypeLike(op, type, roleDescription))) { return failure(); + } auto vecType = cast(type); if (!isa(vecType.getElementType())) return op->emitOpError() @@ -1886,8 +1995,9 @@ static MemoryRole classifyMemoryRole(Type type) { } Attribute memorySpace = memrefType.getMemorySpace(); - if (!memorySpace) + if (!memorySpace) { return MemoryRole::Unknown; + } if (auto addrSpace = dyn_cast(memorySpace)) { switch (addrSpace.getAddressSpace()) { @@ -1934,22 +2044,26 @@ static int64_t getBufferElementByteSize(Type type) { } static Type getBufferElementType(Type type) { - if (auto ptrType = dyn_cast(type)) + if (auto ptrType = dyn_cast(type)) { return ptrType.getElementType(); - if (auto memrefType = dyn_cast(type)) + } + if (auto memrefType = dyn_cast(type)) { return memrefType.getElementType(); + } return {}; } static std::optional getBufferAddressSpace(Type type) { - if (auto ptrType = dyn_cast(type)) + if (auto ptrType = dyn_cast(type)) { return ptrType.getMemorySpace().getAddressSpace(); + } if (auto memrefType = dyn_cast(type)) { if (auto space = dyn_cast_or_null(memrefType.getMemorySpace())) return space.getAddressSpace(); - if (auto intSpace = dyn_cast_or_null(memrefType.getMemorySpace())) + if (auto intSpace = dyn_cast_or_null(memrefType.getMemorySpace())) { return static_cast(intSpace.getInt()); + } } return std::nullopt; } @@ -1962,8 +2076,9 @@ static LogicalResult verifyCubeBridgeLoadLikeOp(BridgeLoadOp op, !isBufferLike(op.getDestination().getType())) return op.emitOpError("requires buffer-like source and destination"); - if (getBufferAddressSpace(op.getSource().getType()) != AddressSpace::MAT) + if (getBufferAddressSpace(op.getSource().getType()) != AddressSpace::MAT) { return op.emitOpError("requires MAT source"); + } if (getBufferAddressSpace(op.getDestination().getType()) != expectedDstSpace) { return op.emitOpError() << "requires " << dstName << " destination"; @@ -1986,8 +2101,9 @@ static LogicalResult verifyCubeBridgeLoadLikeOp(BridgeLoadOp op, static ParseResult parseRequiredOperandWithComma( OpAsmParser &parser, OpAsmParser::UnresolvedOperand &operand) { - if (parser.parseOperand(operand)) + if (parser.parseOperand(operand)) { return failure(); + } (void)parser.parseOptionalComma(); return success(); } @@ -1995,15 +2111,18 @@ static ParseResult parseRequiredOperandWithComma( static ParseResult parseDmaTripleGroup( OpAsmParser &parser, StringRef keyword, SmallVectorImpl &operands) { - if (parser.parseKeyword(keyword) || parser.parseLParen()) + if (parser.parseKeyword(keyword) || parser.parseLParen()) { return failure(); + } for (int i = 0; i < 3; ++i) { OpAsmParser::UnresolvedOperand operand; - if (parser.parseOperand(operand)) + if (parser.parseOperand(operand)) { return failure(); + } operands.push_back(operand); - if (i != 2 && parser.parseComma()) + if (i != 2 && parser.parseComma()) { return failure(); + } } return parser.parseRParen(); } @@ -2014,18 +2133,22 @@ static ParseResult parseOptionalDmaTripleGroupAlias( SmallVectorImpl &operands) { parsedKeyword = {}; for (StringRef keyword : keywords) { - if (failed(parser.parseOptionalKeyword(keyword))) + if (failed(parser.parseOptionalKeyword(keyword))) { continue; + } parsedKeyword = keyword; - if (parser.parseLParen()) + if (parser.parseLParen()) { return failure(); + } for (int i = 0; i < 3; ++i) { OpAsmParser::UnresolvedOperand operand; - if (parser.parseOperand(operand)) + if (parser.parseOperand(operand)) { return failure(); + } operands.push_back(operand); - if (i != 2 && parser.parseComma()) + if (i != 2 && parser.parseComma()) { return failure(); + } } return parser.parseRParen(); } @@ -2033,12 +2156,15 @@ static ParseResult parseOptionalDmaTripleGroupAlias( } static bool isDmaLoopKeyword(StringRef keyword) { - if (keyword == "loop") + if (keyword == "loop") { return true; - if (!keyword.consume_front("loop")) + } + if (!keyword.consume_front("loop")) { return false; - if (keyword.empty()) + } + if (keyword.empty()) { return false; + } return llvm::all_of(keyword, llvm::isDigit); } @@ -2046,11 +2172,13 @@ static ParseResult parseDmaTripleTypes(OpAsmParser &parser, SmallVectorImpl &types) { for (int i = 0; i < 3; ++i) { Type type; - if (parser.parseType(type)) + if (parser.parseType(type)) { return failure(); + } types.push_back(type); - if (i != 2 && parser.parseComma()) + if (i != 2 && parser.parseComma()) { return failure(); + } } return success(); } @@ -2058,8 +2186,9 @@ static ParseResult parseDmaTripleTypes(OpAsmParser &parser, static ParseResult parseDmaPadTypes(OpAsmParser &parser, SmallVectorImpl &types) { Type valueType; - if (parser.parseType(valueType)) + if (parser.parseType(valueType)) { return failure(); + } types.push_back(valueType); if (succeeded(parser.parseOptionalComma())) { Type leftType; @@ -2087,37 +2216,43 @@ static void printDmaTripleTypes(OpAsmPrinter &printer, StringRef keyword, static void printDmaPadGroup(OpAsmPrinter &printer, Value value, Value left, Value right) { printer << " pad(" << value; - if (left || right) + if (left || right) { printer << ", " << left << ", " << right; + } printer << ")"; } static void printDmaPadTypes(OpAsmPrinter &printer, Type valueType, Type leftType, Type rightType) { printer << ", pad " << valueType; - if (leftType || rightType) + if (leftType || rightType) { printer << ", " << leftType << ", " << rightType; + } } static FailureOr parseCubeLoadFracModeKeyword(StringRef keyword) { - if (std::optional mode = symbolizeCubeLoadFracMode(keyword)) + if (std::optional mode = symbolizeCubeLoadFracMode(keyword)) { return *mode; + } return failure(); } static ParseResult parseFixedKeywordOperandGroup( OpAsmParser &parser, StringRef keyword, int operandCount, SmallVectorImpl &operands) { - if (parser.parseKeyword(keyword) || parser.parseLParen()) + if (parser.parseKeyword(keyword) || parser.parseLParen()) { return failure(); + } for (int i = 0; i < operandCount; ++i) { OpAsmParser::UnresolvedOperand operand; - if (parser.parseOperand(operand)) + if (parser.parseOperand(operand)) { return failure(); + } operands.push_back(operand); - if (i + 1 != operandCount && parser.parseComma()) + if (i + 1 != operandCount && parser.parseComma()) { return failure(); + } } return parser.parseRParen(); } @@ -2125,15 +2260,18 @@ static ParseResult parseFixedKeywordOperandGroup( static ParseResult parseFixedKeywordTypes(OpAsmParser &parser, StringRef keyword, int typeCount, SmallVectorImpl &types) { - if (parser.parseKeyword(keyword)) + if (parser.parseKeyword(keyword)) { return failure(); + } for (int i = 0; i < typeCount; ++i) { Type type; - if (parser.parseType(type)) + if (parser.parseType(type)) { return failure(); + } types.push_back(type); - if (i + 1 != typeCount && parser.parseComma()) + if (i + 1 != typeCount && parser.parseComma()) { return failure(); + } } return success(); } @@ -2141,16 +2279,19 @@ static ParseResult parseFixedKeywordTypes(OpAsmParser &parser, StringRef keyword static ParseResult parseCubeLoadFracSrcLayoutGroup( OpAsmParser &parser, SmallVectorImpl &operands) { - if (parser.parseKeyword("src_layout") || parser.parseLParen()) + if (parser.parseKeyword("src_layout") || parser.parseLParen()) { return failure(); + } OpAsmParser::UnresolvedOperand innerStride; - if (parser.parseOperand(innerStride)) + if (parser.parseOperand(innerStride)) { return failure(); + } operands.push_back(innerStride); if (succeeded(parser.parseOptionalComma())) { OpAsmParser::UnresolvedOperand outerStride; - if (parser.parseOperand(outerStride)) + if (parser.parseOperand(outerStride)) { return failure(); + } operands.push_back(outerStride); } return parser.parseRParen(); @@ -2158,16 +2299,19 @@ static ParseResult parseCubeLoadFracSrcLayoutGroup( static ParseResult parseCubeLoadFracSrcLayoutTypes(OpAsmParser &parser, SmallVectorImpl &types) { - if (parser.parseKeyword("src_layout") || parser.parseLParen()) + if (parser.parseKeyword("src_layout") || parser.parseLParen()) { return failure(); + } Type innerStrideType; - if (parser.parseType(innerStrideType)) + if (parser.parseType(innerStrideType)) { return failure(); + } types.push_back(innerStrideType); if (succeeded(parser.parseOptionalComma())) { Type outerStrideType; - if (parser.parseType(outerStrideType)) + if (parser.parseType(outerStrideType)) { return failure(); + } types.push_back(outerStrideType); } return parser.parseRParen(); @@ -2177,8 +2321,9 @@ static void printCubeLoadFracSrcLayoutGroup(OpAsmPrinter &printer, Value srcInnerStride, Value srcOuterStride) { printer << ", src_layout(" << srcInnerStride; - if (srcOuterStride) + if (srcOuterStride) { printer << ", " << srcOuterStride; + } printer << ")"; } @@ -2186,36 +2331,41 @@ static void printCubeLoadFracSrcLayoutTypes(OpAsmPrinter &printer, Type srcInnerStrideType, Type srcOuterStrideType) { printer << ", src_layout(" << srcInnerStrideType; - if (srcOuterStrideType) + if (srcOuterStrideType) { printer << ", " << srcOuterStrideType; + } printer << ")"; } static FailureOr parseAccStoreModeKeyword(StringRef keyword) { - if (std::optional mode = symbolizeAccStoreMode(keyword)) + if (std::optional mode = symbolizeAccStoreMode(keyword)) { return *mode; + } return failure(); } [[maybe_unused]] static ParseResult parseAccStoreModeGroup( OpAsmParser &parser, StringRef &modeKeyword, SmallVectorImpl &modeOperands) { - if (parser.parseKeyword(&modeKeyword)) + if (parser.parseKeyword(&modeKeyword)) { return failure(); + } if (failed(parseAccStoreModeKeyword(modeKeyword))) return parser.emitError(parser.getCurrentLocation(), "expected one of 'nz2nd', 'nz2dn', or 'nz2nz'"); auto parseModeOperandWithParens = [&]() -> ParseResult { OpAsmParser::UnresolvedOperand operand; - if (parser.parseLParen() || parser.parseOperand(operand) || parser.parseRParen()) + if (parser.parseLParen() || parser.parseOperand(operand) || parser.parseRParen()) { return failure(); + } modeOperands.push_back(operand); return success(); }; auto parseModeOperandAfterLParen = [&]() -> ParseResult { OpAsmParser::UnresolvedOperand operand; - if (parser.parseOperand(operand) || parser.parseRParen()) + if (parser.parseOperand(operand) || parser.parseRParen()) { return failure(); + } modeOperands.push_back(operand); return success(); }; @@ -2225,17 +2375,21 @@ static FailureOr parseAccStoreModeKeyword(StringRef keyword) { return success(); case AccStoreMode::Nz2dn: (void)parser.parseOptionalComma(); - if (succeeded(parser.parseOptionalKeyword("loop0_src_stride"))) + if (succeeded(parser.parseOptionalKeyword("loop0_src_stride"))) { return parseModeOperandWithParens(); - if (failed(parser.parseOptionalLParen())) + } + if (failed(parser.parseOptionalLParen())) { return success(); + } return parseModeOperandAfterLParen(); case AccStoreMode::Nz2nz: (void)parser.parseOptionalComma(); - if (succeeded(parser.parseOptionalKeyword("split"))) + if (succeeded(parser.parseOptionalKeyword("split"))) { return parseModeOperandWithParens(); - if (failed(parser.parseOptionalLParen())) + } + if (failed(parser.parseOptionalLParen())) { return success(); + } return parseModeOperandAfterLParen(); } return success(); @@ -2244,19 +2398,22 @@ static FailureOr parseAccStoreModeKeyword(StringRef keyword) { [[maybe_unused]] static ParseResult parseAccStoreModeTypes(OpAsmParser &parser, StringRef modeKeyword, SmallVectorImpl &modeTypes) { - if (parser.parseKeyword(modeKeyword)) + if (parser.parseKeyword(modeKeyword)) { return failure(); + } auto parseModeTypeWithParens = [&]() -> ParseResult { Type modeType; - if (parser.parseLParen() || parser.parseType(modeType) || parser.parseRParen()) + if (parser.parseLParen() || parser.parseType(modeType) || parser.parseRParen()) { return failure(); + } modeTypes.push_back(modeType); return success(); }; auto parseModeTypeAfterLParen = [&]() -> ParseResult { Type modeType; - if (parser.parseType(modeType) || parser.parseRParen()) + if (parser.parseType(modeType) || parser.parseRParen()) { return failure(); + } modeTypes.push_back(modeType); return success(); }; @@ -2266,17 +2423,21 @@ parseAccStoreModeTypes(OpAsmParser &parser, StringRef modeKeyword, return success(); case AccStoreMode::Nz2dn: (void)parser.parseOptionalComma(); - if (succeeded(parser.parseOptionalKeyword("loop0_src_stride"))) + if (succeeded(parser.parseOptionalKeyword("loop0_src_stride"))) { return parseModeTypeWithParens(); - if (failed(parser.parseOptionalLParen())) + } + if (failed(parser.parseOptionalLParen())) { return success(); + } return parseModeTypeAfterLParen(); case AccStoreMode::Nz2nz: (void)parser.parseOptionalComma(); - if (succeeded(parser.parseOptionalKeyword("split"))) + if (succeeded(parser.parseOptionalKeyword("split"))) { return parseModeTypeWithParens(); - if (failed(parser.parseOptionalLParen())) + } + if (failed(parser.parseOptionalLParen())) { return success(); + } return parseModeTypeAfterLParen(); } return success(); @@ -2291,12 +2452,14 @@ parseAccStoreModeTypes(OpAsmParser &parser, StringRef modeKeyword, case AccStoreMode::Nz2nd: return; case AccStoreMode::Nz2dn: - if (loop0SrcStride) + if (loop0SrcStride) { printer << ", loop0_src_stride(" << loop0SrcStride << ")"; + } return; case AccStoreMode::Nz2nz: - if (split) + if (split) { printer << ", split(" << split << ")"; + } return; } llvm_unreachable("unexpected mte_l0c mode"); @@ -2311,12 +2474,14 @@ parseAccStoreModeTypes(OpAsmParser &parser, StringRef modeKeyword, case AccStoreMode::Nz2nd: return; case AccStoreMode::Nz2dn: - if (loop0SrcStrideType) + if (loop0SrcStrideType) { printer << ", loop0_src_stride(" << loop0SrcStrideType << ")"; + } return; case AccStoreMode::Nz2nz: - if (splitType) + if (splitType) { printer << ", split(" << splitType << ")"; + } return; } llvm_unreachable("unexpected mte_l0c mode"); @@ -2343,27 +2508,32 @@ parseAccStoreModeTypes(OpAsmParser &parser, StringRef modeKeyword, [[maybe_unused]] static ParseResult parseMteL0cL1OptionalFpc( OpAsmParser &parser, SmallVectorImpl &fpcOperands) { - if (failed(parser.parseOptionalKeyword("fpc"))) + if (failed(parser.parseOptionalKeyword("fpc"))) { return success(); - if (parser.parseLParen()) + } + if (parser.parseLParen()) { return failure(); + } OpAsmParser::UnresolvedOperand operand; - if (parser.parseOperand(operand) || parser.parseRParen()) + if (parser.parseOperand(operand) || parser.parseRParen()) { return failure(); + } fpcOperands.push_back(operand); return success(); } [[maybe_unused]] static void printMteL0cL1OptionalFpc(OpAsmPrinter &printer, Value fpc) { - if (fpc) + if (fpc) { printer << ", fpc(" << fpc << ")"; + } } [[maybe_unused]] static void printMteL0cL1OptionalFpcType(OpAsmPrinter &printer, Type fpcType) { - if (fpcType) + if (fpcType) { printer << ", fpc(" << fpcType << ")"; + } } [[maybe_unused]] static ParseResult parseMteL0cL1OptionalLoop3Types( @@ -2372,13 +2542,16 @@ printMteL0cL1OptionalFpcType(OpAsmPrinter &printer, Type fpcType) { SmallVectorImpl &loop3DstStrideTypes, StringRef opName) { if (succeeded(parser.parseOptionalComma())) { StringRef keyword; - if (parser.parseKeyword(&keyword)) + if (parser.parseKeyword(&keyword)) { return failure(); - if (keyword != "loop3") + } + if (keyword != "loop3") { return parser.emitError(parser.getCurrentLocation(), "expected 'loop3'"); + } SmallVector loop3GroupTypes; - if (parseDmaTripleTypes(parser, loop3GroupTypes)) + if (parseDmaTripleTypes(parser, loop3GroupTypes)) { return failure(); + } loop3CountTypes.push_back(loop3GroupTypes[0]); loop3SrcStrideTypes.push_back(loop3GroupTypes[1]); loop3DstStrideTypes.push_back(loop3GroupTypes[2]); @@ -2408,20 +2581,25 @@ printMteL0cL1OptionalFpcType(OpAsmPrinter &printer, Type fpcType) { switch (mode) { case AccStoreMode::Nz2nd: - if (split) + if (split) { return op->emitOpError(nz2ndSplitError); - if (loop0SrcStride) + } + if (loop0SrcStride) { return op->emitOpError(nz2ndLoop0Error); + } return success(); case AccStoreMode::Nz2dn: - if (split) + if (split) { return op->emitOpError(nz2dnSplitError); + } return success(); case AccStoreMode::Nz2nz: - if (loop0SrcStride) + if (loop0SrcStride) { return op->emitOpError(nz2nzLoop0Error); - if (loop3Count) + } + if (loop3Count) { return op->emitOpError(nz2nzLoop3Error); + } return success(); } llvm_unreachable("unexpected mte_l0c mode"); @@ -2516,18 +2694,21 @@ static Type getStructuredAccStoreScalingElementType(Value value) { static bool isStructuredAccStoreClipPayloadForUInt8(Type type) { auto intType = dyn_cast(type); - if (!intType || intType.getWidth() != 16) + if (!intType || intType.getWidth() != 16) { return false; + } return intType.isUnsigned() || intType.isSignless(); } static bool isStructuredAccStoreClipPayloadForSignedInt(Type type) { auto intType = dyn_cast(type); - if (!intType) + if (!intType) { return false; + } unsigned width = intType.getWidth(); - if (width != 4 && width != 8 && width != 16) + if (width != 4 && width != 8 && width != 16) { return false; + } return intType.isSigned() || intType.isSignless(); } @@ -2544,13 +2725,16 @@ static bool isStructuredAccStoreFloatScalarPayload(Value value) { } static bool isStructuredAccStoreClipSupportedElementType(Type type) { - if (auto floatType = dyn_cast(type)) + if (auto floatType = dyn_cast(type)) { return floatType.isF16(); + } auto intType = dyn_cast(type); - if (!intType) + if (!intType) { return false; - if (intType.isUnsignedInteger(8)) + } + if (intType.isUnsignedInteger(8)) { return true; + } if (intType.isSignlessInteger(4) || intType.isSignlessInteger(8) || intType.isSignlessInteger(16)) return true; @@ -2563,13 +2747,15 @@ static bool isStructuredAccStoreClipSupportedElementType(Type type) { static LogicalResult verifyStructuredAccStoreClipPayload(Operation *op, Type destinationElementType, Value clipValue) { - if (!clipValue) + if (!clipValue) { return success(); + } Type clipType = clipValue.getType(); if (destinationElementType.isF16()) { - if (!clipType.isF16()) + if (!clipType.isF16()) { return op->emitOpError("clip for f16 destination requires f16 payload"); + } return success(); } @@ -2580,16 +2766,18 @@ static LogicalResult verifyStructuredAccStoreClipPayload(Operation *op, << destinationElementType; if (intType.isUnsignedInteger(8)) { - if (!isStructuredAccStoreClipPayloadForUInt8(clipType)) + if (!isStructuredAccStoreClipPayloadForUInt8(clipType)) { return op->emitOpError("clip for ui8 destination requires ui16/signless i16 payload"); + } return success(); } if (intType.isSignlessInteger(4) || intType.isSignlessInteger(8) || intType.isSignlessInteger(16) || intType.isSignedInteger(4) || intType.isSignedInteger(8) || intType.isSignedInteger(16)) { - if (!isStructuredAccStoreClipPayloadForSignedInt(clipType)) + if (!isStructuredAccStoreClipPayloadForSignedInt(clipType)) { return op->emitOpError("clip for signed 4/8/16-bit destination requires signed/signless i4/i8/i16 payload"); + } return success(); } @@ -2725,20 +2913,24 @@ static bool isStructuredAccStoreDestinationFamily( case StructuredAccStoreDestinationFamily::BF16: return type.isBF16(); case StructuredAccStoreDestinationFamily::I32: - if (auto intType = dyn_cast(type)) + if (auto intType = dyn_cast(type)) { return intType.getWidth() == 32; + } return false; case StructuredAccStoreDestinationFamily::I16: - if (auto intType = dyn_cast(type)) + if (auto intType = dyn_cast(type)) { return intType.getWidth() == 16 && !intType.isUnsigned(); + } return false; case StructuredAccStoreDestinationFamily::I8: - if (auto intType = dyn_cast(type)) + if (auto intType = dyn_cast(type)) { return intType.getWidth() == 8; + } return false; case StructuredAccStoreDestinationFamily::I4: - if (auto intType = dyn_cast(type)) + if (auto intType = dyn_cast(type)) { return intType.getWidth() == 4 && !intType.isUnsigned(); + } return false; case StructuredAccStoreDestinationFamily::FP8: return pto::isPTOFloat8Type(type) || pto::isPTOHiFloat8Type(type) || @@ -2749,15 +2941,19 @@ static bool isStructuredAccStoreDestinationFamily( static ParseResult parseStructuredAccStoreUnitFlag(OpAsmParser &parser, StructuredAccStoreAsmState &state) { - if (state.unitFlag) + if (state.unitFlag) { return parser.emitError(parser.getCurrentLocation(), "duplicate unit_flag clause"); + } StringRef keyword; - if (parser.parseLParen() || parser.parseKeyword(&keyword) || parser.parseRParen()) + if (parser.parseLParen() || parser.parseKeyword(&keyword) || parser.parseRParen()) { return failure(); - if (keyword == "check_only") + } + if (keyword == "check_only") { state.unitFlag = AccStoreUnitFlagCtrl::CheckOnly; - else if (keyword == "check_and_clear") + } + else if (keyword == "check_and_clear") { state.unitFlag = AccStoreUnitFlagCtrl::CheckAndClear; + } else return parser.emitError(parser.getCurrentLocation(), "expected 'check_only' or 'check_and_clear'"); @@ -2766,8 +2962,9 @@ static ParseResult parseStructuredAccStoreUnitFlag(OpAsmParser &parser, static ParseResult parseStructuredAccStorePreQuant( OpAsmParser &parser, StructuredAccStoreAsmState &state) { - if (state.preQuantMode) + if (state.preQuantMode) { return parser.emitError(parser.getCurrentLocation(), "duplicate pre_quant clause"); + } OpAsmParser::UnresolvedOperand payload; StringRef modeKeyword; if (parser.parseLParen() || parser.parseOperand(payload) || parser.parseComma() || @@ -2775,8 +2972,9 @@ static ParseResult parseStructuredAccStorePreQuant( parser.parseKeyword(&modeKeyword) || parser.parseRParen()) return failure(); auto mode = symbolizeAccStoreQuantPreMode(modeKeyword); - if (!mode) + if (!mode) { return parser.emitError(parser.getCurrentLocation(), "invalid pre_quant mode"); + } state.preQuantOperands.push_back(payload); state.preQuantMode = *mode; return success(); @@ -2784,40 +2982,48 @@ static ParseResult parseStructuredAccStorePreQuant( static ParseResult parseStructuredAccStorePreRelu( OpAsmParser &parser, StructuredAccStoreAsmState &state) { - if (state.preReluMode) + if (state.preReluMode) { return parser.emitError(parser.getCurrentLocation(), "duplicate pre_relu clause"); + } StringRef modeKeyword; bool hasPayload = false; OpAsmParser::UnresolvedOperand payload; - if (parser.parseLParen()) + if (parser.parseLParen()) { return failure(); + } if (failed(parser.parseOptionalKeyword("mode"))) { hasPayload = true; if (parser.parseOperand(payload) || parser.parseComma() || parser.parseKeyword("mode")) return failure(); } - if (parser.parseEqual() || parser.parseKeyword(&modeKeyword)) + if (parser.parseEqual() || parser.parseKeyword(&modeKeyword)) { return failure(); + } auto mode = symbolizeReluPreMode(modeKeyword); - if (!mode) + if (!mode) { return parser.emitError(parser.getCurrentLocation(), "invalid pre_relu mode"); + } if (succeeded(parser.parseOptionalComma())) { - if (parser.parseKeyword("clip") || parser.parseEqual()) + if (parser.parseKeyword("clip") || parser.parseEqual()) { return failure(); + } if (!state.clipValueOperands.empty()) return parser.emitError(parser.getCurrentLocation(), "duplicate clip payload in pre_relu clause"); OpAsmParser::UnresolvedOperand clipValue; - if (parser.parseOperand(clipValue)) + if (parser.parseOperand(clipValue)) { return failure(); + } state.clipValueOperands.push_back(clipValue); } - if (parser.parseRParen()) + if (parser.parseRParen()) { return failure(); + } - if (hasPayload) + if (hasPayload) { state.preReluOperands.push_back(payload); + } state.preReluMode = *mode; return success(); } @@ -2828,21 +3034,24 @@ static ParseResult parseStructuredAccStoreLayout( if (failed(mode)) return parser.emitError(parser.getCurrentLocation(), "expected one of 'nz2nd', 'nz2dn', or 'nz2nz'"); - if (state.mode) + if (state.mode) { return parser.emitError(parser.getCurrentLocation(), "duplicate layout clause"); + } state.mode = *mode; if (*mode == AccStoreMode::Nz2dn) { if (succeeded(parser.parseOptionalLParen())) { OpAsmParser::UnresolvedOperand operand; - if (parser.parseOperand(operand) || parser.parseRParen()) + if (parser.parseOperand(operand) || parser.parseRParen()) { return failure(); + } state.loop0SrcStrideOperands.push_back(operand); } } else if (*mode == AccStoreMode::Nz2nz) { if (succeeded(parser.parseOptionalLParen())) { OpAsmParser::UnresolvedOperand operand; - if (parser.parseOperand(operand) || parser.parseRParen()) + if (parser.parseOperand(operand) || parser.parseRParen()) { return failure(); + } state.splitOperands.push_back(operand); } } @@ -2851,8 +3060,9 @@ static ParseResult parseStructuredAccStoreLayout( static ParseResult parseStructuredAccStoreLoop3( OpAsmParser &parser, StructuredAccStoreAsmState &state) { - if (!state.loop3CountOperands.empty()) + if (!state.loop3CountOperands.empty()) { return parser.emitError(parser.getCurrentLocation(), "duplicate loop3 clause"); + } OpAsmParser::UnresolvedOperand count; OpAsmParser::UnresolvedOperand srcStride; OpAsmParser::UnresolvedOperand dstStride; @@ -2868,8 +3078,9 @@ static ParseResult parseStructuredAccStoreLoop3( static ParseResult parseStructuredAccStoreAtomic( OpAsmParser &parser, StructuredAccStoreAsmState &state) { - if (state.atomicType || state.atomicOp) + if (state.atomicType || state.atomicOp) { return parser.emitError(parser.getCurrentLocation(), "duplicate atomic clause"); + } StringRef typeKeyword; StringRef opKeyword; if (parser.parseLParen() || parser.parseKeyword("type") || parser.parseEqual() || @@ -2879,10 +3090,12 @@ static ParseResult parseStructuredAccStoreAtomic( return failure(); auto type = symbolizeAccStoreAtomicType(typeKeyword); auto op = symbolizeAccStoreAtomicOp(opKeyword); - if (!type) + if (!type) { return parser.emitError(parser.getCurrentLocation(), "invalid atomic type"); - if (!op) + } + if (!op) { return parser.emitError(parser.getCurrentLocation(), "invalid atomic op"); + } state.atomicType = *type; state.atomicOp = *op; return success(); @@ -2894,33 +3107,42 @@ static ParseResult parseStructuredAccStoreClauses( bool seenClause = false; while (true) { if (seenClause) { - if (failed(parser.parseOptionalComma())) + if (failed(parser.parseOptionalComma())) { return success(); + } } StringRef keyword; OptionalParseResult optParseResult = parser.parseOptionalKeyword(&keyword); if (!optParseResult.has_value() || failed(*optParseResult)) { - if (!seenClause) + if (!seenClause) { return success(); + } return parser.emitError(parser.getCurrentLocation(), "expected valid keyword"); } seenClause = true; StructuredAccStoreClauseKind kind; - if (keyword == "unit_flag") + if (keyword == "unit_flag") { kind = StructuredAccStoreClauseKind::UnitFlag; - else if (keyword == "pre_quant") + } + else if (keyword == "pre_quant") { kind = StructuredAccStoreClauseKind::PreQuant; - else if (keyword == "pre_relu") + } + else if (keyword == "pre_relu") { kind = StructuredAccStoreClauseKind::PreRelu; - else if (keyword == "nz2nd" || keyword == "nz2dn" || keyword == "nz2nz") + } + else if (keyword == "nz2nd" || keyword == "nz2dn" || keyword == "nz2nz") { kind = StructuredAccStoreClauseKind::Layout; - else if (keyword == "loop3") + } + else if (keyword == "loop3") { kind = StructuredAccStoreClauseKind::Loop3; - else if (keyword == "sat" || keyword == "nosat") + } + else if (keyword == "sat" || keyword == "nosat") { kind = StructuredAccStoreClauseKind::Sat; - else if (keyword == "atomic") + } + else if (keyword == "atomic") { kind = StructuredAccStoreClauseKind::Atomic; + } else return parser.emitError(parser.getCurrentLocation(), "unknown mte_l0c clause"); @@ -2948,8 +3170,9 @@ static ParseResult parseStructuredAccStoreClauses( parseResult = parseStructuredAccStoreLoop3(parser, state); break; case StructuredAccStoreClauseKind::Sat: - if (state.satMode) + if (state.satMode) { return parser.emitError(parser.getCurrentLocation(), "duplicate sat/nosat clause"); + } if (keyword == "nosat") { state.satMode = AccStoreSatMode::NoSat; break; @@ -2959,8 +3182,9 @@ static ParseResult parseStructuredAccStoreClauses( if (parser.parseKeyword(&satOption) || satOption != "preserve_nan") return parser.emitError(parser.getCurrentLocation(), "expected preserve_nan"); - if (parser.parseRParen()) + if (parser.parseRParen()) { return failure(); + } state.satMode = AccStoreSatMode::SatPreserveNan; } else { state.satMode = AccStoreSatMode::Sat; @@ -2970,16 +3194,18 @@ static ParseResult parseStructuredAccStoreClauses( parseResult = parseStructuredAccStoreAtomic(parser, state); break; } - if (failed(parseResult)) + if (failed(parseResult)) { return failure(); + } } } static ParseResult parseStructuredOptionalType(OpAsmParser &parser, SmallVectorImpl &types) { Type type; - if (parser.parseType(type)) + if (parser.parseType(type)) { return failure(); + } types.push_back(type); return success(); } @@ -2995,25 +3221,29 @@ static LogicalResult verifyStructuredAccStoreLike( std::optional atomicType, std::optional atomicOp, bool allowAtomic) { auto getBufferElementType = [](Type type) -> Type { - if (auto ptrType = dyn_cast(type)) + if (auto ptrType = dyn_cast(type)) { return ptrType.getElementType(); - if (auto memrefType = dyn_cast(type)) + } + if (auto memrefType = dyn_cast(type)) { return memrefType.getElementType(); + } return {}; }; Type sourceElementType = getBufferElementType(srcType); Type destinationElementType = getBufferElementType(dstType); - if (static_cast(preQuant) != static_cast(preQuantMode)) + if (static_cast(preQuant) != static_cast(preQuantMode)) { return op->emitOpError("pre_quant requires payload and mode together"); + } if (preQuantMode) { if (*preQuantMode == AccStoreQuantPreMode::NoConvert) { // The no_convert keyword carries no quantization parameters; the // syntactic payload operand is ignored for compatibility with the // structured pre_quant clause form. } else if (isStructuredAccStoreVectorQuantMode(*preQuantMode)) { - if (!isStructuredAccStoreScalingPayload(preQuant)) + if (!isStructuredAccStoreScalingPayload(preQuant)) { return op->emitOpError("vector pre_quant mode requires scaling pointer payload"); + } if (!isStructuredAccStoreFloatScalarPayloadType( getStructuredAccStoreScalingElementType(preQuant))) return op->emitOpError( @@ -3032,11 +3262,13 @@ static LogicalResult verifyStructuredAccStoreLike( if (*preQuantMode != AccStoreQuantPreMode::NoConvert) { if (isa(sourceElementType)) { - if (!isStructuredAccStoreFloatPreQuantMode(*preQuantMode)) + if (!isStructuredAccStoreFloatPreQuantMode(*preQuantMode)) { return emitIncompatibleQuantModeError(); + } } else if (sourceElementType.isSignlessInteger(32)) { - if (!isStructuredAccStoreInt32PreQuantMode(*preQuantMode)) + if (!isStructuredAccStoreInt32PreQuantMode(*preQuantMode)) { return emitIncompatibleQuantModeError(); + } } else { return op->emitOpError() << "pre_quant requires source element type to be f32 or i32, got " @@ -3060,31 +3292,39 @@ static LogicalResult verifyStructuredAccStoreLike( return failure(); if (!preReluMode) { - if (preRelu) + if (preRelu) { return op->emitOpError("pre_relu payload requires pre_relu mode"); - if (clipValue) + } + if (clipValue) { return op->emitOpError("clip requires pre_relu clause"); + } } else { switch (*preReluMode) { case ReluPreMode::NoRelu: - if (preRelu) + if (preRelu) { return op->emitOpError("mode does not accept pre_relu payload"); + } break; case ReluPreMode::NormalRelu: - if (preRelu) + if (preRelu) { return op->emitOpError("mode does not accept pre_relu payload"); + } break; case ReluPreMode::ScalarRelu: - if (!preRelu) + if (!preRelu) { return op->emitOpError("scalar_relu requires payload"); - if (!isStructuredAccStoreFloatScalarPayload(preRelu)) + } + if (!isStructuredAccStoreFloatScalarPayload(preRelu)) { return op->emitOpError("scalar_relu requires f16/bf16/f32 payload"); + } break; case ReluPreMode::VectorRelu: - if (!preRelu) + if (!preRelu) { return op->emitOpError("vector_relu requires payload"); - if (!isStructuredAccStoreScalingPayload(preRelu)) + } + if (!isStructuredAccStoreScalingPayload(preRelu)) { return op->emitOpError("vector_relu requires scaling pointer payload"); + } if (!isStructuredAccStoreFloatScalarPayloadType( getStructuredAccStoreScalingElementType(preRelu))) return op->emitOpError( @@ -3097,29 +3337,37 @@ static LogicalResult verifyStructuredAccStoreLike( bool hasLoop3 = static_cast(loop3Count) || static_cast(loop3SrcStride) || static_cast(loop3DstStride); - if (hasLoop3 && !(loop3Count && loop3SrcStride && loop3DstStride)) + if (hasLoop3 && !(loop3Count && loop3SrcStride && loop3DstStride)) { return op->emitOpError("loop3 requires count, src stride, and dst stride together"); + } if (!mode) { - if (split) + if (split) { return op->emitOpError("split requires nz2nz"); - if (loop0SrcStride) + } + if (loop0SrcStride) { return op->emitOpError("loop0_src_stride requires nz2dn"); - if (loop3Count) + } + if (loop3Count) { return op->emitOpError("loop3 requires nz2nd or nz2dn"); + } } else { switch (*mode) { case AccStoreMode::Nz2nd: - if (split) + if (split) { return op->emitOpError("nz2nd does not accept split"); - if (loop0SrcStride) + } + if (loop0SrcStride) { return op->emitOpError("nz2nd does not accept loop0_src_stride"); + } break; case AccStoreMode::Nz2dn: { - if (!loop0SrcStride) + if (!loop0SrcStride) { return op->emitOpError("nz2dn requires loop0_src_stride"); - if (split) + } + if (split) { return op->emitOpError("nz2dn does not accept split"); + } APInt loop0Value; if (unitFlag && *unitFlag != AccStoreUnitFlagCtrl::Off && (!matchPattern(loop0SrcStride, m_ConstantInt(&loop0Value)) || @@ -3130,10 +3378,12 @@ static LogicalResult verifyStructuredAccStoreLike( break; } case AccStoreMode::Nz2nz: - if (loop0SrcStride) + if (loop0SrcStride) { return op->emitOpError("nz2nz does not accept loop0_src_stride"); - if (loop3Count) + } + if (loop3Count) { return op->emitOpError("loop3 requires nz2nd or nz2dn"); + } if (!isa(destinationElementType) || !cast(destinationElementType).isF32()) return op->emitOpError("nz2nz requires destination element type to be f32"); @@ -3141,10 +3391,12 @@ static LogicalResult verifyStructuredAccStoreLike( } } - if (static_cast(atomicType) != static_cast(atomicOp)) + if (static_cast(atomicType) != static_cast(atomicOp)) { return op->emitOpError("atomic requires type and op together"); - if ((atomicType || atomicOp) && !allowAtomic) + } + if ((atomicType || atomicOp) && !allowAtomic) { return op->emitOpError("atomic is only supported for mte_l0c_gm"); + } return success(); } @@ -3171,11 +3423,13 @@ static void printStructuredAccStoreClauses( } if (preReluMode) { printer << ", pre_relu("; - if (preRelu) + if (preRelu) { printer << preRelu << ", "; + } printer << "mode = " << stringifyReluPreMode(*preReluMode); - if (clipValue) + if (clipValue) { printer << ", clip = " << clipValue; + } printer << ")"; } if (mode) { @@ -3185,13 +3439,15 @@ static void printStructuredAccStoreClauses( break; case AccStoreMode::Nz2dn: printer << ", nz2dn"; - if (loop0SrcStride) + if (loop0SrcStride) { printer << "(" << loop0SrcStride << ")"; + } break; case AccStoreMode::Nz2nz: printer << ", nz2nz"; - if (split) + if (split) { printer << "(" << split << ")"; + } break; } } @@ -3222,16 +3478,21 @@ static void printStructuredAccStoreOptionalTypes( OpAsmPrinter &printer, Value preQuant, Value preRelu, Value clipValue, Value split, Value loop0SrcStride, Value loop3Count, Value loop3SrcStride, Value loop3DstStride) { - if (preQuant) + if (preQuant) { printer << ", " << preQuant.getType(); - if (preRelu) + } + if (preRelu) { printer << ", " << preRelu.getType(); - if (clipValue) + } + if (clipValue) { printer << ", " << clipValue.getType(); - if (split) + } + if (split) { printer << ", " << split.getType(); - if (loop0SrcStride) + } + if (loop0SrcStride) { printer << ", " << loop0SrcStride.getType(); + } if (loop3Count) printer << ", " << loop3Count.getType() << ", " << loop3SrcStride.getType() << ", " << loop3DstStride.getType(); @@ -3389,10 +3650,12 @@ static LogicalResult verifyCopyGmToUbufOp(CopyOp op, bool expectSourceGM) { int64_t sourceElemBytes = getBufferElementByteSize(op.getSource().getType()); int64_t destinationElemBytes = getBufferElementByteSize(op.getDestination().getType()); - if (sourceElemBytes <= 0 || destinationElemBytes <= 0) + if (sourceElemBytes <= 0 || destinationElemBytes <= 0) { return op.emitOpError("requires copy source and destination element types with known byte width"); - if (sourceElemBytes != destinationElemBytes) + } + if (sourceElemBytes != destinationElemBytes) { return op.emitOpError("requires source and destination element byte widths to match"); + } return success(); } @@ -3451,10 +3714,12 @@ static LogicalResult verifyCopyUbufToGmOp(CopyOp op, bool expectSourceGM) { int64_t sourceElemBytes = getBufferElementByteSize(op.getSource().getType()); int64_t destinationElemBytes = getBufferElementByteSize(op.getDestination().getType()); - if (sourceElemBytes <= 0 || destinationElemBytes <= 0) + if (sourceElemBytes <= 0 || destinationElemBytes <= 0) { return op.emitOpError("requires copy source and destination element types with known byte width"); - if (sourceElemBytes != destinationElemBytes) + } + if (sourceElemBytes != destinationElemBytes) { return op.emitOpError("requires source and destination element byte widths to match"); + } return success(); } @@ -3473,10 +3738,12 @@ static LogicalResult verifyCopyCbufToUbufLikeOp(CopyOp op) { int64_t sourceElemBytes = getBufferElementByteSize(op.getSource().getType()); int64_t destinationElemBytes = getBufferElementByteSize(op.getDestination().getType()); - if (sourceElemBytes <= 0 || destinationElemBytes <= 0) + if (sourceElemBytes <= 0 || destinationElemBytes <= 0) { return op.emitOpError("requires copy source and destination element types with known byte width"); - if (sourceElemBytes != destinationElemBytes) + } + if (sourceElemBytes != destinationElemBytes) { return op.emitOpError("requires source and destination element byte widths to match"); + } return success(); } @@ -3532,8 +3799,9 @@ LogicalResult VRegType::verify(function_ref emitError, LogicalResult VecScopeOp::verify() { Region &bodyRegion = getBody(); - if (bodyRegion.empty()) + if (bodyRegion.empty()) { return emitOpError("expects a non-empty body region"); + } Block &body = bodyRegion.front(); if (body.getNumArguments() != 0) @@ -3545,8 +3813,9 @@ LogicalResult VecScopeOp::verify() { LogicalResult StrictVecScopeOp::verify() { Region &bodyRegion = getBody(); - if (bodyRegion.empty()) + if (bodyRegion.empty()) { return emitOpError("expects a non-empty body region"); + } Block &body = bodyRegion.front(); if (body.getNumArguments() != getCaptures().size()) @@ -3613,12 +3882,15 @@ void MteGmUbOp::build(OpBuilder &builder, OperationState &state, Value source, std::optional pad) { state.addOperands({source, destination, l2CacheCtl, lenBurst, nburst.count, nburst.srcStride, nburst.dstStride}); - for (const pto::DmaLoopConfig &loop : loops) + for (const pto::DmaLoopConfig &loop : loops) { state.addOperands(loop.count); - for (const pto::DmaLoopConfig &loop : loops) + } + for (const pto::DmaLoopConfig &loop : loops) { state.addOperands(loop.srcStride); - for (const pto::DmaLoopConfig &loop : loops) + } + for (const pto::DmaLoopConfig &loop : loops) { state.addOperands(loop.dstStride); + } bool hasPadCounts = pad && pad->leftCount && pad->rightCount; assert((!pad || static_cast(pad->leftCount) == static_cast(pad->rightCount)) && @@ -3646,10 +3918,12 @@ void MteGmUbOp::build(OpBuilder &builder, OperationState &state, Value source, std::optional loop2, std::optional pad) { SmallVector loops; - if (loop1) + if (loop1) { loops.push_back(*loop1); - if (loop2) + } + if (loop2) { loops.push_back(*loop2); + } build(builder, state, source, destination, l2CacheCtl, lenBurst, nburst, loops, pad); } @@ -3669,11 +3943,13 @@ ParseResult MteGmUbOp::parse(OpAsmParser &parser, OperationState &result) { return failure(); while (true) { if (succeeded(parser.parseOptionalKeyword("pad"))) { - if (parser.parseLParen()) + if (parser.parseLParen()) { return failure(); + } OpAsmParser::UnresolvedOperand value; - if (parser.parseOperand(value)) + if (parser.parseOperand(value)) { return failure(); + } padOperands.push_back(value); if (succeeded(parser.parseOptionalComma())) { OpAsmParser::UnresolvedOperand left; @@ -3684,8 +3960,9 @@ ParseResult MteGmUbOp::parse(OpAsmParser &parser, OperationState &result) { padOperands.push_back(left); padOperands.push_back(right); } - if (parser.parseRParen()) + if (parser.parseRParen()) { return failure(); + } break; } @@ -3694,15 +3971,17 @@ ParseResult MteGmUbOp::parse(OpAsmParser &parser, OperationState &result) { if (parseOptionalDmaTripleGroupAlias(parser, {"loop", "loop1", "loop2"}, parsedKeyword, loopGroupOperands)) return failure(); - if (parsedKeyword.empty()) + if (parsedKeyword.empty()) { break; + } loopCountOperands.push_back(loopGroupOperands[0]); loopSrcStrideOperands.push_back(loopGroupOperands[1]); loopDstStrideOperands.push_back(loopGroupOperands[2]); } - if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) + if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) { return failure(); + } Type sourceType, destinationType, l2CacheCtlType, lenBurstType; SmallVector nburstTypes, loopCountTypes, loopSrcStrideTypes, @@ -3715,20 +3994,23 @@ ParseResult MteGmUbOp::parse(OpAsmParser &parser, OperationState &result) { return failure(); while (succeeded(parser.parseOptionalComma())) { StringRef keyword; - if (parser.parseKeyword(&keyword)) + if (parser.parseKeyword(&keyword)) { return failure(); + } if (isDmaLoopKeyword(keyword)) { SmallVector loopGroupTypes; - if (parseDmaTripleTypes(parser, loopGroupTypes)) + if (parseDmaTripleTypes(parser, loopGroupTypes)) { return failure(); + } loopCountTypes.push_back(loopGroupTypes[0]); loopSrcStrideTypes.push_back(loopGroupTypes[1]); loopDstStrideTypes.push_back(loopGroupTypes[2]); continue; } if (keyword == "pad") { - if (!padTypes.empty() || parseDmaPadTypes(parser, padTypes)) + if (!padTypes.empty() || parseDmaPadTypes(parser, padTypes)) { return failure(); + } continue; } return parser.emitError(parser.getCurrentLocation(), @@ -3809,14 +4091,16 @@ void MteGmUbOp::getEffects( } LogicalResult MteGmUbOp::verify() { - if (failed(verifyCopyGmToUbufOp(*this, true))) + if (failed(verifyCopyGmToUbufOp(*this, true))) { return failure(); + } if (failed(verifyDmaLoadStoreLoopGroups( getOperation(), getLoopCounts(), getLoopSrcStrides(), getLoopDstStrides()))) return failure(); - if (!getPadValue() && (getLeftPaddingCount() || getRightPaddingCount())) + if (!getPadValue() && (getLeftPaddingCount() || getRightPaddingCount())) { return emitOpError() << "requires pad group to provide a pad value"; + } if (getPadValue() && static_cast(getLeftPaddingCount()) != static_cast(getRightPaddingCount())) return emitOpError() @@ -3833,8 +4117,9 @@ LogicalResult MteGmUbOp::verify() { LogicalResult SetMovPadValOp::verify() { Type valueType = getValue().getType(); - if (isSupportedMovPadScalarType(valueType)) + if (isSupportedMovPadScalarType(valueType)) { return success(); + } return emitOpError() << "expects i8/i16/i32 or f16/bf16/f32 scalar operand, but got " << valueType; @@ -3854,8 +4139,9 @@ static LogicalResult verifyMadPointerKinds(Operation *op, Type lhsTy, Type rhsTy auto lhsType = dyn_cast(lhsTy); auto rhsType = dyn_cast(rhsTy); auto dstType = dyn_cast(dstTy); - if (!lhsType || !rhsType || !dstType) + if (!lhsType || !rhsType || !dstType) { return op->emitOpError("requires typed !pto.ptr lhs/rhs/dst operands"); + } const auto lhsAS = lhsType.getMemorySpace().getAddressSpace(); const auto rhsAS = rhsType.getMemorySpace().getAddressSpace(); @@ -3864,15 +4150,18 @@ static LogicalResult verifyMadPointerKinds(Operation *op, Type lhsTy, Type rhsTy const bool isStrongCube = lhsAS == pto::AddressSpace::LEFT && rhsAS == pto::AddressSpace::RIGHT && dstAS == pto::AddressSpace::ACC; - if (!isStrongCube) + if (!isStrongCube) { return op->emitOpError("requires l0a/l0b/l0c-typed lhs/rhs/dst pointers"); + } - if (!biasTy) + if (!biasTy) { return success(); + } auto biasType = dyn_cast(*biasTy); - if (!biasType) + if (!biasType) { return op->emitOpError("requires typed !pto.ptr bias operand"); + } if (biasType.getMemorySpace().getAddressSpace() != pto::AddressSpace::BIAS) { return op->emitOpError("requires bias pointer in !pto.ptr<..., bt>"); } @@ -3905,8 +4194,9 @@ static LogicalResult verifyMadMxCommon(Operation *op, Type lhsTy, Type rhsTy, Type dstTy, std::optional biasTy = std::nullopt) { - if (failed(verifyMadPointerKinds(op, lhsTy, rhsTy, dstTy, biasTy))) + if (failed(verifyMadPointerKinds(op, lhsTy, rhsTy, dstTy, biasTy))) { return failure(); + } auto lhsType = cast(lhsTy); auto rhsType = cast(rhsTy); @@ -3917,8 +4207,9 @@ static LogicalResult verifyMadMxCommon(Operation *op, Type lhsTy, Type rhsTy, const bool isStrongCube = lhsAS == pto::AddressSpace::LEFT && rhsAS == pto::AddressSpace::RIGHT && dstAS == pto::AddressSpace::ACC; - if (!isStrongCube) + if (!isStrongCube) { return op->emitOpError("requires l0a/l0b/l0c-typed lhs/rhs/dst pointers"); + } if (!isMxElementType(lhsType.getElementType()) || !isMxElementType(rhsType.getElementType())) { @@ -3959,10 +4250,12 @@ void MadMxBiasOp::getEffects( static std::optional parseMadUnitFlagModeToken(StringRef token) { - if (token == "check_only") + if (token == "check_only") { return pto::MadUnitFlagMode::CheckOnly; - if (token == "check_and_set") + } + if (token == "check_and_set") { return pto::MadUnitFlagMode::CheckAndSet; + } return std::nullopt; } @@ -3977,10 +4270,12 @@ static StringRef stringifyMadUnitFlagModeToken(pto::MadUnitFlagMode mode) { } static std::optional parseTf32ModeToken(StringRef token) { - if (token == "round_even") + if (token == "round_even") { return pto::Tf32Mode::RoundEven; - if (token == "round_away") + } + if (token == "round_away") { return pto::Tf32Mode::RoundAway; + } return std::nullopt; } @@ -4010,14 +4305,16 @@ static LogicalResult verifyMadSemanticClauses(Operation *op, Type lhsTy, std::optional tf32Mode, std::optional satMode, bool hasNDir) { - if (failed(verifyMadPointerKinds(op, lhsTy, rhsTy, dstTy, biasTy))) + if (failed(verifyMadPointerKinds(op, lhsTy, rhsTy, dstTy, biasTy))) { return failure(); + } auto lhsType = dyn_cast(lhsTy); auto rhsType = dyn_cast(rhsTy); auto dstType = dyn_cast(dstTy); - if (!lhsType || !rhsType || !dstType) + if (!lhsType || !rhsType || !dstType) { return op->emitOpError("requires typed !pto.ptr lhs/rhs/dst operands"); + } if (tf32Mode) { if (!(lhsType.getElementType().isF32() && rhsType.getElementType().isF32() && @@ -4033,8 +4330,9 @@ static LogicalResult verifyMadSemanticClauses(Operation *op, Type lhsTy, } if (satMode) { auto isFloatLike = [](Type type) { - if (isa(type)) + if (isa(type)) { return true; + } return pto::isPTOLowPrecisionType(type); }; if (!(isFloatLike(lhsType.getElementType()) && @@ -4069,8 +4367,9 @@ static ParseResult parseMadSemanticOpCommon(OpAsmParser &parser, return failure(); auto parseUnitFlagClause = [&]() -> ParseResult { - if (failed(parser.parseOptionalKeyword("unit_flag"))) + if (failed(parser.parseOptionalKeyword("unit_flag"))) { return success(); + } if (parser.parseLParen() || parser.parseKeyword(&unitFlagKeyword) || parser.parseRParen()) return failure(); @@ -4103,10 +4402,12 @@ static ParseResult parseMadSemanticOpCommon(OpAsmParser &parser, return success(); }; auto parseTf32Clause = [&]() -> ParseResult { - if (!parseTf32ModeClause) + if (!parseTf32ModeClause) { return success(); - if (failed(parser.parseOptionalKeyword("tf32_mode"))) + } + if (failed(parser.parseOptionalKeyword("tf32_mode"))) { return success(); + } if (parser.parseLParen() || parser.parseKeyword(&tf32Keyword) || parser.parseRParen()) return failure(); @@ -4129,8 +4430,9 @@ static ParseResult parseMadSemanticOpCommon(OpAsmParser &parser, failed(parseNDirClause())) return failure(); - if (parser.parseOptionalAttrDict(attrs) || parser.parseColon()) + if (parser.parseOptionalAttrDict(attrs) || parser.parseColon()) { return failure(); + } Type lhsType, rhsType, dstType, mType, nType, kType, biasType; if (parser.parseType(lhsType) || parser.parseComma() || @@ -4138,8 +4440,9 @@ static ParseResult parseMadSemanticOpCommon(OpAsmParser &parser, parser.parseType(dstType) || parser.parseComma()) return failure(); if (hasBias) { - if (parser.parseType(biasType) || parser.parseComma()) + if (parser.parseType(biasType) || parser.parseComma()) { return failure(); + } } if (parser.parseType(mType) || parser.parseComma() || parser.parseType(nType) || parser.parseComma() || parser.parseType(kType)) @@ -4174,18 +4477,21 @@ static void printMadSemanticClauses(OpAsmPrinter &printer, Operation *op, printer << " unit_flag(" << stringifyMadUnitFlagModeToken(unitFlagMode.getValue()) << ")"; } - if (op->hasAttr("disable_gemv")) + if (op->hasAttr("disable_gemv")) { printer << " disable_gemv"; - if (auto satMode = op->getAttrOfType("sat_mode")) + } + if (auto satMode = op->getAttrOfType("sat_mode")) { printer << ' ' << stringifyMadSatModeToken(satMode.getValue()); + } if (allowTf32Mode) { if (auto tf32Mode = op->getAttrOfType("tf32_mode")) { printer << " tf32_mode(" << stringifyTf32ModeToken(tf32Mode.getValue()) << ")"; } } - if (op->hasAttr("n_dir")) + if (op->hasAttr("n_dir")) { printer << " n_dir"; + } } static ArrayRef getMadSemanticElidedAttrs(bool allowTf32Mode) { @@ -4454,32 +4760,39 @@ Value MadMxBiasRawOp::getBiasOrNull() { return getBias(); } static bool isCompatibleScalarForSemanticType(Type semanticType, Type scalarType) { - if (semanticType == scalarType) + if (semanticType == scalarType) { return true; + } auto semanticInt = dyn_cast(semanticType); auto scalarInt = dyn_cast(scalarType); - if (!semanticInt || !scalarInt || semanticInt.getWidth() != scalarInt.getWidth()) + if (!semanticInt || !scalarInt || semanticInt.getWidth() != scalarInt.getWidth()) { return false; + } - if (semanticInt.isSigned()) + if (semanticInt.isSigned()) { return scalarInt.isSigned() || scalarInt.isSignless(); - if (semanticInt.isUnsigned()) + } + if (semanticInt.isUnsigned()) { return scalarInt.isUnsigned() || scalarInt.isSignless(); + } return scalarInt.isSignless(); } LogicalResult VbrOp::verify() { - if (failed(verifyVRegTypeLike(*this, getResult().getType(), "result"))) + if (failed(verifyVRegTypeLike(*this, getResult().getType(), "result"))) { return failure(); + } auto resultVecType = cast(getResult().getType()); Type elementType = getValue().getType(); - if (isa(elementType)) + if (isa(elementType)) { return emitOpError("value must be a scalar matching the result element type"); + } Type resultElementType = resultVecType.getElementType(); - if (!isCompatibleScalarForSemanticType(resultElementType, elementType)) + if (!isCompatibleScalarForSemanticType(resultElementType, elementType)) { return emitOpError("value type must match result element type"); + } return success(); } @@ -4492,8 +4805,9 @@ static LogicalResult verifyWideningReductionVecOp(ReductionOp op, auto inputType = dyn_cast(op.getInput().getType()); auto resultType = dyn_cast(op.getResult().getType()); - if (!inputType || !resultType) + if (!inputType || !resultType) { return failure(); + } Type inputElemType = inputType.getElementType(); Type expectedResultElemType = inputElemType; @@ -4534,8 +4848,9 @@ LogicalResult VcmaxOp::verify() { if (failed(verifyVRegTypeLike(*this, getInput().getType(), "input")) || failed(verifyVRegTypeLike(*this, getResult().getType(), "result"))) return failure(); - if (getInput().getType() != getResult().getType()) + if (getInput().getType() != getResult().getType()) { return emitOpError("input and result must have the same vector type"); + } return success(); } @@ -4543,25 +4858,29 @@ LogicalResult VcminOp::verify() { if (failed(verifyVRegTypeLike(*this, getInput().getType(), "input")) || failed(verifyVRegTypeLike(*this, getResult().getType(), "result"))) return failure(); - if (getInput().getType() != getResult().getType()) + if (getInput().getType() != getResult().getType()) { return emitOpError("input and result must have the same vector type"); + } return success(); } LogicalResult VciOp::verify() { auto resultType = dyn_cast(getResult().getType()); - if (!resultType) + if (!resultType) { return emitOpError("result must be !pto.vreg<...>"); + } Type resultElemType = resultType.getElementType(); bool supportedInteger = false; if (auto intType = dyn_cast(resultElemType)) supportedInteger = intType.getWidth() == 8 || intType.getWidth() == 16 || intType.getWidth() == 32; bool supportedFloat = resultElemType.isF16() || resultElemType.isF32(); - if (!supportedInteger && !supportedFloat) + if (!supportedInteger && !supportedFloat) { return emitOpError("result element type must be integer or f16/f32"); - if (!isCompatibleScalarForSemanticType(resultElemType, getIndex().getType())) + } + if (!isCompatibleScalarForSemanticType(resultElemType, getIndex().getType())) { return emitOpError("index type must match result element type"); + } return success(); } @@ -4581,8 +4900,9 @@ static bool isSameVgather2IntegerSemantics(IntegerType sourceType, if (!sourceType || !resultType || sourceType.getWidth() != resultType.getWidth()) return false; - if (sourceType.isUnsigned()) + if (sourceType.isUnsigned()) { return resultType.isUnsigned(); + } return !resultType.isUnsigned(); } @@ -4591,28 +4911,34 @@ static bool isVgather2B8ResultType(IntegerType sourceType, if (!sourceType || !resultType || sourceType.getWidth() != 8 || resultType.getWidth() != 16) return false; - if (sourceType.isUnsigned()) + if (sourceType.isUnsigned()) { return resultType.isUnsigned(); + } return !resultType.isUnsigned(); } LogicalResult Vgather2Op::verify() { - if (!isBufferLike(getSource().getType())) + if (!isBufferLike(getSource().getType())) { return emitOpError("requires a pointer-like source"); + } MemoryRole sourceRole = classifyMemoryRole(getSource().getType()); - if (sourceRole == MemoryRole::GM) + if (sourceRole == MemoryRole::GM) { return emitOpError("requires a UB-backed source"); + } auto offsetsType = dyn_cast(getOffsets().getType()); auto resultType = dyn_cast(getResult().getType()); - if (!offsetsType || !resultType) + if (!offsetsType || !resultType) { return emitOpError("offsets and result must be !pto.vreg<...>"); + } auto offsetsElemType = dyn_cast(offsetsType.getElementType()); - if (!offsetsElemType) + if (!offsetsElemType) { return emitOpError("offset vector must use integer element type"); - if (offsetsType.getElementCount() != resultType.getElementCount()) + } + if (offsetsType.getElementCount() != resultType.getElementCount()) { return emitOpError("offset and result vectors must have the same element count"); + } Type sourceElemType = getBufferElementType(getSource().getType()); Type resultElemType = resultType.getElementType(); @@ -4662,8 +4988,9 @@ LogicalResult Vgather2Op::verify() { "requires source element type i8/ui8/i16/ui16/i32/ui32/f16/bf16/f32"); } - if (resultElemWidth != 16 && resultElemWidth != 32) + if (resultElemWidth != 16 && resultElemWidth != 32) { return emitOpError("result element type must be 16-bit or 32-bit"); + } if (resultType.getElementCount() != expectedLanes) return emitOpError() << "expects result type " << formatVRegType(expectedLanes, resultElemType); @@ -4678,8 +5005,9 @@ LogicalResult Vgather2Op::verify() { } LogicalResult CopyUbufToUbufOp::verify() { - if (!isBufferLike(getSource().getType()) || !isBufferLike(getDestination().getType())) + if (!isBufferLike(getSource().getType()) || !isBufferLike(getDestination().getType())) { return emitOpError("requires pointer-like source and destination"); + } if (classifyMemoryRole(getSource().getType()) != MemoryRole::UB || classifyMemoryRole(getDestination().getType()) != MemoryRole::UB) return emitOpError("requires UB-backed source and destination"); @@ -4705,8 +5033,9 @@ void CopyUbufToCbufOp::getEffects( } LogicalResult CopyUbufToCbufOp::verify() { - if (!isBufferLike(getSource().getType()) || !isBufferLike(getDestination().getType())) + if (!isBufferLike(getSource().getType()) || !isBufferLike(getDestination().getType())) { return emitOpError("requires pointer-like source and destination"); + } if (classifyMemoryRole(getSource().getType()) != MemoryRole::UB || classifyMemoryRole(getDestination().getType()) != MemoryRole::Other) return emitOpError("requires UB-backed source and CBUF-backed destination"); @@ -4721,8 +5050,9 @@ void MteUbUbOp::getEffects( } LogicalResult MteUbUbOp::verify() { - if (!isBufferLike(getSource().getType()) || !isBufferLike(getDestination().getType())) + if (!isBufferLike(getSource().getType()) || !isBufferLike(getDestination().getType())) { return emitOpError("requires pointer-like source and destination"); + } if (classifyMemoryRole(getSource().getType()) != MemoryRole::UB || classifyMemoryRole(getDestination().getType()) != MemoryRole::UB) return emitOpError("requires UB-backed source and destination"); @@ -4737,8 +5067,9 @@ void MteUbL1Op::getEffects( } LogicalResult MteUbL1Op::verify() { - if (!isBufferLike(getSource().getType()) || !isBufferLike(getDestination().getType())) + if (!isBufferLike(getSource().getType()) || !isBufferLike(getDestination().getType())) { return emitOpError("requires pointer-like source and destination"); + } if (classifyMemoryRole(getSource().getType()) != MemoryRole::UB || classifyMemoryRole(getDestination().getType()) != MemoryRole::Other) return emitOpError("requires UB-backed source and CBUF-backed destination"); @@ -4752,11 +5083,13 @@ void VgatherbOp::getEffects( } LogicalResult VgatherbOp::verify() { - if (!isBufferLike(getSource().getType())) + if (!isBufferLike(getSource().getType())) { return emitOpError("requires a pointer-like source"); + } MemoryRole sourceRole = classifyMemoryRole(getSource().getType()); - if (sourceRole == MemoryRole::GM) + if (sourceRole == MemoryRole::GM) { return emitOpError("requires a UB-backed source"); + } if (failed(verifyMaskTypeWithGranularityLike(getOperation(), getMask().getType(), "mask type", "b32"))) @@ -4764,13 +5097,16 @@ LogicalResult VgatherbOp::verify() { auto offsetsType = dyn_cast(getOffsets().getType()); auto resultType = dyn_cast(getResult().getType()); - if (!offsetsType || !resultType) + if (!offsetsType || !resultType) { return emitOpError("offsets and result must be !pto.vreg<...>"); + } auto offsetsElemType = dyn_cast(offsetsType.getElementType()); - if (!offsetsElemType) + if (!offsetsElemType) { return emitOpError("offset vector must use integer element type"); - if (offsetsElemType.getWidth() != 32) + } + if (offsetsElemType.getWidth() != 32) { return emitOpError("currently requires 32-bit offset vector elements"); + } // vgatherb is a 32-byte block gather: each offset addresses one 32-byte block. // The offset vector holds VL/32 block addresses (always ui32), while the // result vector holds VL/sizeof(T) elements of the data type. These counts @@ -4788,24 +5124,31 @@ void Vgather2BcOp::getEffects( } LogicalResult Vgather2BcOp::verify() { - if (!isBufferLike(getSource().getType())) + if (!isBufferLike(getSource().getType())) { return emitOpError("requires a pointer-like source"); - if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed source"); - if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) + } + if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) { return failure(); + } auto offsetsType = dyn_cast(getOffsets().getType()); auto resultType = dyn_cast(getResult().getType()); - if (!offsetsType || !resultType) + if (!offsetsType || !resultType) { return emitOpError("offsets and result must be !pto.vreg<...>"); + } auto offsetsElemType = dyn_cast(offsetsType.getElementType()); - if (!offsetsElemType) + if (!offsetsElemType) { return emitOpError("offset vector must use integer element type"); - if (offsetsElemType.getWidth() != 32) + } + if (offsetsElemType.getWidth() != 32) { return emitOpError("currently requires 32-bit offset vector elements"); - if (offsetsType.getElementCount() != resultType.getElementCount()) + } + if (offsetsType.getElementCount() != resultType.getElementCount()) { return emitOpError("offset and result vectors must have the same element count"); + } return success(); } @@ -4817,10 +5160,12 @@ LogicalResult VbitsortOp::verify() { classifyMemoryRole(getSource().getType()) != MemoryRole::UB || classifyMemoryRole(getIndices().getType()) != MemoryRole::UB) return emitOpError("requires UB-backed destination/source/indices"); - if (!getRepeatTimes().getType().isIndex()) + if (!getRepeatTimes().getType().isIndex()) { return emitOpError("repeat_times must be index"); - if (failed(verifyNotNestedInVecScope(*this, "pto.vbitsort"))) + } + if (failed(verifyNotNestedInVecScope(*this, "pto.vbitsort"))) { return failure(); + } return success(); } @@ -4859,10 +5204,12 @@ LogicalResult Vmrgsort4Op::verify() { src3PtrType.getElementType() != elemType) return emitOpError( "requires destination and all sources to have the same element type"); - if (!elemType.isF16() && !elemType.isF32()) + if (!elemType.isF16() && !elemType.isF32()) { return emitOpError("requires f16 or f32 element type"); - if (failed(verifyNotNestedInVecScope(*this, "pto.vmrgsort4"))) + } + if (failed(verifyNotNestedInVecScope(*this, "pto.vmrgsort4"))) { return failure(); + } return success(); } @@ -4896,15 +5243,18 @@ void VldsOp::getEffects( template static LogicalResult verifyVldsCommon(LoadOp op) { - if (!isBufferLike(op.getSource().getType())) + if (!isBufferLike(op.getSource().getType())) { return op.emitOpError("requires a pointer-like source"); + } - if (failed(verifyVRegTypeLike(op, op.getResult().getType(), "result type"))) + if (failed(verifyVRegTypeLike(op, op.getResult().getType(), "result type"))) { return failure(); + } MemoryRole sourceRole = classifyMemoryRole(op.getSource().getType()); - if (sourceRole == MemoryRole::GM) + if (sourceRole == MemoryRole::GM) { return op.emitOpError("requires a UB-backed source"); + } if (op.getDistAttr()) { StringRef dist = *op.getDist(); @@ -4919,11 +5269,13 @@ static LogicalResult verifyVldsCommon(LoadOp op) { } LogicalResult VldsOp::verify() { - if (failed(verifyVldsCommon(*this))) + if (failed(verifyVldsCommon(*this))) { return failure(); + } if (Value updatedBase = getUpdatedBase()) { - if (updatedBase.getType() != getSource().getType()) + if (updatedBase.getType() != getSource().getType()) { return emitOpError("requires updated base result to match base type"); + } } return success(); } @@ -4934,12 +5286,15 @@ void VldasOp::getEffects( } LogicalResult VldasOp::verify() { - if (!isBufferLike(getSource().getType())) + if (!isBufferLike(getSource().getType())) { return emitOpError("requires a pointer-like source"); - if (failed(verifyAlignTypeLike(*this, getResult().getType(), "result type"))) + } + if (failed(verifyAlignTypeLike(*this, getResult().getType(), "result type"))) { return failure(); - if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed source"); + } return success(); } @@ -4948,10 +5303,12 @@ LogicalResult InitAlignOp::verify() { } LogicalResult SprclrOp::verify() { - if (!isSupportedSprToken(getSpr())) + if (!isSupportedSprToken(getSpr())) { return emitOpError("requires spr to be \"AR\""); - if (failed(verifyNestedInVecScope(*this, "pto.sprclr"))) + } + if (failed(verifyNestedInVecScope(*this, "pto.sprclr"))) { return failure(); + } return success(); } @@ -4959,27 +5316,35 @@ static LogicalResult verifySprStoreCommon(Operation *op, StringRef opName, StringRef spr, Value destination, Value offset, bool requireImmediateOffset) { - if (!isSupportedSprToken(spr)) + if (!isSupportedSprToken(spr)) { return op->emitOpError("requires spr to be \"AR\""); - if (failed(verifyNestedInVecScope(op, opName))) + } + if (failed(verifyNestedInVecScope(op, opName))) { return failure(); + } auto ptrType = dyn_cast(destination.getType()); - if (!ptrType) + if (!ptrType) { return op->emitOpError("requires a pointer-like UB destination"); - if (classifyMemoryRole(destination.getType()) != MemoryRole::UB) + } + if (classifyMemoryRole(destination.getType()) != MemoryRole::UB) { return op->emitOpError("requires a UB-backed destination"); + } auto intType = dyn_cast(ptrType.getElementType()); - if (!intType || intType.getWidth() != 32 || intType.isSigned()) + if (!intType || intType.getWidth() != 32 || intType.isSigned()) { return op->emitOpError("requires ui32/i32 UB destination element type"); - if (!offset.getType().isInteger(32)) + } + if (!offset.getType().isInteger(32)) { return op->emitOpError("requires i32 offset"); + } if (requireImmediateOffset) { APInt offsetValue; - if (!matchPattern(offset, m_ConstantInt(&offsetValue))) + if (!matchPattern(offset, m_ConstantInt(&offsetValue))) { return op->emitOpError("requires constant immediate offset"); + } int64_t signedOffset = offsetValue.getSExtValue(); - if (signedOffset < -128 || signedOffset > 127) + if (signedOffset < -128 || signedOffset > 127) { return op->emitOpError("requires signed 8-bit immediate offset"); + } } return success(); } @@ -5030,15 +5395,18 @@ LogicalResult VldusOp::verify() { failed(verifyAlignTypeLike(*this, getUpdatedAlign().getType(), "updated align type"))) return failure(); - if (!isBufferLike(getSource().getType())) + if (!isBufferLike(getSource().getType())) { return emitOpError("requires a pointer-like source"); - if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed source"); + } if (static_cast(getIncrement()) != static_cast(getUpdatedBase())) return emitOpError( "requires increment and updated base result to appear together"); - if (getUpdatedBase() && getUpdatedBase().getType() != getSource().getType()) + if (getUpdatedBase() && getUpdatedBase().getType() != getSource().getType()) { return emitOpError("requires updated base result to match source type"); + } return success(); } @@ -5049,16 +5417,20 @@ void UvldOp::getEffects( } LogicalResult UvldOp::verify() { - if (failed(verifyVRegTypeLike(*this, getResult().getType(), "result type"))) + if (failed(verifyVRegTypeLike(*this, getResult().getType(), "result type"))) { return failure(); - if (!isBufferLike(getSource().getType())) + } + if (!isBufferLike(getSource().getType())) { return emitOpError("requires a buffer-like source"); - if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed source"); + } auto sourceMemRef = dyn_cast(getSource().getType()); - if (!sourceMemRef) + if (!sourceMemRef) { return success(); + } Type sourceElementType = sourceMemRef.getElementType(); Type vectorElementType = cast(getResult().getType()).getElementType(); @@ -5070,33 +5442,39 @@ LogicalResult UvldOp::verify() { LogicalResult VdupOp::verify() { auto resultType = dyn_cast(getResult().getType()); - if (!resultType) + if (!resultType) { return emitOpError("result must be !pto.vreg<...>"); + } std::optional granularity = getVdupMaskGranularity(resultType.getElementType()); - if (!granularity) + if (!granularity) { return emitOpError("result element type must use b8, b16, or b32 mask granularity"); + } if (failed(verifyMaskTypeWithGranularityLike( getOperation(), getMask().getType(), "mask type", *granularity))) return failure(); - if (!isSupportedVdupPosition(getPosition())) + if (!isSupportedVdupPosition(getPosition())) { return emitOpError("position must be LOWEST or HIGHEST"); + } Type inputType = getInput().getType(); if (auto inputVecType = dyn_cast(inputType)) { - if (inputVecType != resultType) + if (inputVecType != resultType) { return emitOpError("vector input must match result vector type"); + } return success(); } - if (getPosition()) + if (getPosition()) { return emitOpError("position is only supported for vector input"); + } Type resultElementType = resultType.getElementType(); - if (!isCompatibleScalarForSemanticType(resultElementType, inputType)) + if (!isCompatibleScalarForSemanticType(resultElementType, inputType)) { return emitOpError("scalar input must match result element type"); + } return success(); } @@ -5120,8 +5498,9 @@ LogicalResult TensorViewAddrOp::verify() { expectedRank = memrefType.getRank(); auto srcSpace = dyn_cast_or_null(memrefType.getMemorySpace()); - if (srcSpace && srcSpace != gmSpace) + if (srcSpace && srcSpace != gmSpace) { return emitOpError("memref source must stay in gm memory space"); + } } else { return emitOpError( "source must be a tensor_view, partition_tensor_view, or memref"); @@ -5131,23 +5510,27 @@ LogicalResult TensorViewAddrOp::verify() { if (dstMemRefType.getElementType() != elementType) return emitOpError( "memref result element type must match source element type"); - if (dstMemRefType.getRank() != expectedRank) + if (dstMemRefType.getRank() != expectedRank) { return emitOpError("memref result rank must match source rank"); + } auto dstSpace = dyn_cast_or_null(dstMemRefType.getMemorySpace()); - if (dstSpace && dstSpace != gmSpace) + if (dstSpace && dstSpace != gmSpace) { return emitOpError("memref result must stay in gm memory space"); + } return success(); } auto dstPtrType = dyn_cast(dstType); - if (!dstPtrType) + if (!dstPtrType) { return emitOpError("result must be a memref or !pto.ptr<...>"); + } if (dstPtrType.getElementType() != elementType) return emitOpError( "pointer result element type must match source element type"); - if (dstPtrType.getMemorySpace() != gmSpace) + if (dstPtrType.getMemorySpace() != gmSpace) { return emitOpError("pointer result must stay in gm memory space"); + } return success(); } @@ -5175,23 +5558,27 @@ LogicalResult TileBufAddrOp::verify() { if (dstMemRefType.getElementType() != elementType) return emitOpError( "memref result element type must match tile element type"); - if (dstMemRefType.getRank() != srcRank) + if (dstMemRefType.getRank() != srcRank) { return emitOpError("memref result rank must match tile rank"); + } auto dstSpace = dyn_cast_or_null(dstMemRefType.getMemorySpace()); - if (srcSpace && dstSpace && srcSpace != dstSpace) + if (srcSpace && dstSpace && srcSpace != dstSpace) { return emitOpError("memref result must stay within the tile memory space"); + } return success(); } auto dstPtrType = dyn_cast(dstType); - if (!dstPtrType) + if (!dstPtrType) { return emitOpError("result must be a memref or !pto.ptr<...>"); + } if (dstPtrType.getElementType() != elementType) return emitOpError( "pointer result element type must match tile element type"); - if (srcSpace && dstPtrType.getMemorySpace() != srcSpace) + if (srcSpace && dstPtrType.getMemorySpace() != srcSpace) { return emitOpError("pointer result must stay within the tile memory space"); + } return success(); } @@ -5200,8 +5587,9 @@ LogicalResult PsetB8Op::verify() { "result type", "b8"))) return failure(); - if (!isSupportedPredicatePattern(getPattern())) + if (!isSupportedPredicatePattern(getPattern())) { return emitOpError("requires a supported PAT_* predicate pattern"); + } return success(); } @@ -5210,8 +5598,9 @@ LogicalResult PsetB16Op::verify() { "result type", "b16"))) return failure(); - if (!isSupportedPredicatePattern(getPattern())) + if (!isSupportedPredicatePattern(getPattern())) { return emitOpError("requires a supported PAT_* predicate pattern"); + } return success(); } @@ -5219,8 +5608,9 @@ LogicalResult PsetB32Op::verify() { if (failed(verifyMaskTypeWithGranularityLike(*this, getResult().getType(), "result type", "b32"))) return failure(); - if (!isSupportedPredicatePattern(getPattern())) + if (!isSupportedPredicatePattern(getPattern())) { return emitOpError("requires a supported PAT_* predicate pattern"); + } return success(); } @@ -5228,8 +5618,9 @@ LogicalResult PgeB8Op::verify() { if (failed(verifyMaskTypeWithGranularityLike(*this, getResult().getType(), "result type", "b8"))) return failure(); - if (!isSupportedPredicatePattern(getPattern())) + if (!isSupportedPredicatePattern(getPattern())) { return emitOpError("requires a supported PAT_* predicate pattern"); + } return success(); } @@ -5237,8 +5628,9 @@ LogicalResult PgeB16Op::verify() { if (failed(verifyMaskTypeWithGranularityLike(*this, getResult().getType(), "result type", "b16"))) return failure(); - if (!isSupportedPredicatePattern(getPattern())) + if (!isSupportedPredicatePattern(getPattern())) { return emitOpError("requires a supported PAT_* predicate pattern"); + } return success(); } @@ -5246,8 +5638,9 @@ LogicalResult PgeB32Op::verify() { if (failed(verifyMaskTypeWithGranularityLike(*this, getResult().getType(), "result type", "b32"))) return failure(); - if (!isSupportedPredicatePattern(getPattern())) + if (!isSupportedPredicatePattern(getPattern())) { return emitOpError("requires a supported PAT_* predicate pattern"); + } return success(); } @@ -5259,10 +5652,12 @@ static LogicalResult verifyPredicateLaneCountOp(PltOp op, return failure(); Type scalarType = op.getScalar().getType(); auto scalarIntType = dyn_cast(scalarType); - if (!scalarIntType || scalarIntType.getWidth() != 32) + if (!scalarIntType || scalarIntType.getWidth() != 32) { return op.emitOpError("requires scalar to be i32"); - if (op.getScalarOut().getType() != scalarType) + } + if (op.getScalarOut().getType() != scalarType) { return op.emitOpError("requires scalar_out to match scalar type"); + } return success(); } @@ -5280,10 +5675,12 @@ static LogicalResult verifyPredicateLoopBoundOp(PltmOp op, if (failed(verifyMaskTypeWithGranularityLike(op, op.getMask().getType(), "mask type", granularity))) return failure(); - if (!op.getLoop().getType().isInteger(16)) + if (!op.getLoop().getType().isInteger(16)) { return op.emitOpError("requires loop operand to be i16"); - if (!op.getBound().getType().isInteger(32)) + } + if (!op.getBound().getType().isInteger(32)) { return op.emitOpError("requires bound operand to be i32"); + } return success(); } @@ -5301,8 +5698,9 @@ LogicalResult PpackOp::verify() { if (failed(verifyMaskTypeLike(*this, getInput().getType(), "input type")) || failed(verifyMaskTypeLike(*this, getResult().getType(), "result type"))) return failure(); - if (!isSupportedPartToken(getPart())) + if (!isSupportedPartToken(getPart())) { return emitOpError("requires part to be LOWER or HIGHER"); + } auto inputMaskType = cast(getInput().getType()); auto resultMaskType = cast(getResult().getType()); StringRef inputGranularity = inputMaskType.getGranularity(); @@ -5319,8 +5717,9 @@ LogicalResult PunpackOp::verify() { if (failed(verifyMaskTypeLike(*this, getInput().getType(), "input type")) || failed(verifyMaskTypeLike(*this, getResult().getType(), "result type"))) return failure(); - if (!isSupportedPartToken(getPart())) + if (!isSupportedPartToken(getPart())) { return emitOpError("requires part to be LOWER or HIGHER"); + } auto inputMaskType = cast(getInput().getType()); auto resultMaskType = cast(getResult().getType()); StringRef inputGranularity = inputMaskType.getGranularity(); @@ -5378,17 +5777,22 @@ void PldsOp::getEffects( } LogicalResult PldsOp::verify() { - if (!isBufferLike(getSource().getType())) + if (!isBufferLike(getSource().getType())) { return emitOpError("requires a pointer-like source"); - if (failed(verifyMaskTypeLike(*this, getResult().getType(), "result type"))) + } + if (failed(verifyMaskTypeLike(*this, getResult().getType(), "result type"))) { return failure(); + } MemoryRole sourceRole = classifyMemoryRole(getSource().getType()); - if (sourceRole == MemoryRole::GM) + if (sourceRole == MemoryRole::GM) { return emitOpError("requires a UB-backed source"); - if (!getOffset().getType().isIndex()) + } + if (!getOffset().getType().isIndex()) { return emitOpError("requires index offset"); - if (!isSupportedPredicateLoadDist(getDist())) + } + if (!isSupportedPredicateLoadDist(getDist())) { return emitOpError("requires predicate load dist to be NORM, US, or DS"); + } if (getUpdatedBase() && getUpdatedBase().getType() != getSource().getType()) return emitOpError("requires updated base result to match base type"); @@ -5402,16 +5806,21 @@ void PldiOp::getEffects( } LogicalResult PldiOp::verify() { - if (!isBufferLike(getSource().getType())) + if (!isBufferLike(getSource().getType())) { return emitOpError("requires a pointer-like source"); - if (failed(verifyMaskTypeLike(*this, getResult().getType(), "result type"))) + } + if (failed(verifyMaskTypeLike(*this, getResult().getType(), "result type"))) { return failure(); - if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed source"); - if (!matchPattern(getOffset(), m_Constant())) + } + if (!matchPattern(getOffset(), m_Constant())) { return emitOpError("requires offset to be a constant index immediate"); - if (!isSupportedPredicateLoadDist(getDist())) + } + if (!isSupportedPredicateLoadDist(getDist())) { return emitOpError("requires predicate load dist to be NORM, US, or DS"); + } if (getUpdatedBase() && getUpdatedBase().getType() != getSource().getType()) return emitOpError("requires updated base result to match base type"); @@ -5422,28 +5831,34 @@ template static LogicalResult verifyElementwiseVecScalarOpLike(OpTy op) { auto inputType = dyn_cast(op.getInput().getType()); auto resultType = dyn_cast(op.getResult().getType()); - if (!inputType || !resultType) + if (!inputType || !resultType) { return op.emitOpError("input and result must be !pto.vreg<...>"); - if (inputType != resultType) + } + if (inputType != resultType) { return op.emitOpError("input and result vector types must match"); + } Type elemType = inputType.getElementType(); Type scalarType = op.getScalar().getType(); - if (scalarType == elemType) + if (scalarType == elemType) { return success(); + } auto elemInt = dyn_cast(elemType); auto scalarInt = dyn_cast(scalarType); - if (!elemInt || !scalarInt || elemInt.getWidth() != scalarInt.getWidth()) + if (!elemInt || !scalarInt || elemInt.getWidth() != scalarInt.getWidth()) { return op.emitOpError("scalar type must match vector element type"); + } - if (elemInt.isSigned() && (scalarInt.isSigned() || scalarInt.isSignless())) + if (elemInt.isSigned() && (scalarInt.isSigned() || scalarInt.isSignless())) { return success(); + } if (elemInt.isUnsigned() && (scalarInt.isUnsigned() || scalarInt.isSignless())) return success(); - if (elemInt.isSignless() && scalarInt.isSignless()) + if (elemInt.isSignless() && scalarInt.isSignless()) { return success(); + } return op.emitOpError( "integer scalar type must match vector element width and use matching signedness or signless i"); @@ -5451,17 +5866,20 @@ static LogicalResult verifyElementwiseVecScalarOpLike(OpTy op) { template static LogicalResult verifyVecScalarOpLike(OpTy op) { - if (failed(verifyElementwiseVecScalarOpLike(op))) + if (failed(verifyElementwiseVecScalarOpLike(op))) { return failure(); + } return success(); } template static LogicalResult verifyVecScalarMaskedOpLike(OpTy op) { - if (failed(verifyElementwiseVecScalarOpLike(op))) + if (failed(verifyElementwiseVecScalarOpLike(op))) { return failure(); - if (failed(verifyMaskTypeLike(op, op.getMask().getType(), "mask type"))) + } + if (failed(verifyMaskTypeLike(op, op.getMask().getType(), "mask type"))) { return failure(); + } if (failed(verifyNonLowPrecisionVRegElementTypeLike( op.getOperation(), op.getInput().getType(), "input type"))) return failure(); @@ -5482,10 +5900,12 @@ static LogicalResult verifyCarryVecOp(CarryOp op) { auto rhsType = cast(op.getRhs().getType()); auto resultType = cast(op.getResult().getType()); auto lhsElemType = cast(lhsType.getElementType()); - if (lhsType != rhsType || lhsType != resultType) + if (lhsType != rhsType || lhsType != resultType) { return op.emitOpError("requires lhs, rhs, and result to have matching vector types"); - if (lhsElemType.getWidth() != 32) + } + if (lhsElemType.getWidth() != 32) { return op.emitOpError("currently requires 32-bit integer vector elements"); + } return success(); } @@ -5506,61 +5926,79 @@ LogicalResult VlreluOp::verify() { return verifyVecScalarMaskedOpLike(*this); } LogicalResult VshlsOp::verify() { auto inputType = dyn_cast(getInput().getType()); auto resultType = dyn_cast(getResult().getType()); - if (!inputType || !resultType) + if (!inputType || !resultType) { return emitOpError("input and result must be !pto.vreg<...>"); - if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) + } + if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) { return failure(); - if (inputType != resultType) + } + if (inputType != resultType) { return emitOpError("input and result vector types must match"); - if (!isa(inputType.getElementType())) + } + if (!isa(inputType.getElementType())) { return emitOpError("requires integer vector and integer scalar"); + } auto scalarType = dyn_cast(getScalar().getType()); - if (!scalarType || !scalarType.isSignlessInteger(16)) + if (!scalarType || !scalarType.isSignlessInteger(16)) { return emitOpError("requires signless i16 scalar"); + } return success(); } LogicalResult VshrsOp::verify() { auto inputType = dyn_cast(getInput().getType()); auto resultType = dyn_cast(getResult().getType()); - if (!inputType || !resultType) + if (!inputType || !resultType) { return emitOpError("input and result must be !pto.vreg<...>"); - if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) + } + if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) { return failure(); - if (inputType != resultType) + } + if (inputType != resultType) { return emitOpError("input and result vector types must match"); - if (!isa(inputType.getElementType())) + } + if (!isa(inputType.getElementType())) { return emitOpError("requires integer vector and integer scalar"); + } auto scalarType = dyn_cast(getScalar().getType()); - if (!scalarType || !scalarType.isSignlessInteger(16)) + if (!scalarType || !scalarType.isSignlessInteger(16)) { return emitOpError("requires signless i16 scalar"); + } return success(); } LogicalResult VabsOp::verify() { - if (failed(verifyVRegTypeLike(*this, getInput().getType(), "operand type"))) + if (failed(verifyVRegTypeLike(*this, getInput().getType(), "operand type"))) { return failure(); - if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) + } + if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) { return failure(); - if (failed(verifyVRegTypeLike(*this, getResult().getType(), "result type"))) + } + if (failed(verifyVRegTypeLike(*this, getResult().getType(), "result type"))) { return failure(); - if (getInput().getType() != getResult().getType()) + } + if (getInput().getType() != getResult().getType()) { return emitOpError("requires matching register vector shape"); + } return success(); } template static LogicalResult verifyUnaryVecOp(UnaryOp op) { - if (failed(verifyVRegTypeLike(op, op.getInput().getType(), "operand type"))) + if (failed(verifyVRegTypeLike(op, op.getInput().getType(), "operand type"))) { return failure(); - if (failed(verifyMaskTypeLike(op, op.getMask().getType(), "mask type"))) + } + if (failed(verifyMaskTypeLike(op, op.getMask().getType(), "mask type"))) { return failure(); - if (failed(verifyVRegTypeLike(op, op.getResult().getType(), "result type"))) + } + if (failed(verifyVRegTypeLike(op, op.getResult().getType(), "result type"))) { return failure(); + } if (failed(verifyNonLowPrecisionVRegElementTypeLike( op.getOperation(), op.getInput().getType(), "operand type"))) return failure(); - if (op.getInput().getType() != op.getResult().getType()) + if (op.getInput().getType() != op.getResult().getType()) { return op.emitOpError("requires matching register vector shape"); + } return success(); } @@ -5569,17 +6007,20 @@ LogicalResult VlnOp::verify() { return verifyUnaryVecOp(*this); } LogicalResult VsqrtOp::verify() { return verifyUnaryVecOp(*this); } LogicalResult VnegOp::verify() { return verifyUnaryVecOp(*this); } LogicalResult VreluOp::verify() { - if (failed(verifyUnaryVecOp(*this))) + if (failed(verifyUnaryVecOp(*this))) { return failure(); + } auto inputType = cast(getInput().getType()); Type elemType = inputType.getElementType(); if (auto intType = dyn_cast(elemType)) { - if (intType.getWidth() != 32 || intType.isUnsigned()) + if (intType.getWidth() != 32 || intType.isUnsigned()) { return emitOpError("requires si32/i32/f16/f32 vector element type"); + } return success(); } - if (!elemType.isF16() && !elemType.isF32()) + if (!elemType.isF16() && !elemType.isF32()) { return emitOpError("requires si32/i32/f16/f32 vector element type"); + } return success(); } LogicalResult VnotOp::verify() { return verifyUnaryVecOp(*this); } @@ -5587,14 +6028,18 @@ LogicalResult VnotOp::verify() { return verifyUnaryVecOp(*this); } template static LogicalResult verifyBinaryVecOp(BinaryOp op, bool allowLowPrecision = false) { - if (failed(verifyVRegTypeLike(op, op.getLhs().getType(), "lhs type"))) + if (failed(verifyVRegTypeLike(op, op.getLhs().getType(), "lhs type"))) { return failure(); - if (failed(verifyVRegTypeLike(op, op.getRhs().getType(), "rhs type"))) + } + if (failed(verifyVRegTypeLike(op, op.getRhs().getType(), "rhs type"))) { return failure(); - if (failed(verifyMaskTypeLike(op, op.getMask().getType(), "mask type"))) + } + if (failed(verifyMaskTypeLike(op, op.getMask().getType(), "mask type"))) { return failure(); - if (failed(verifyVRegTypeLike(op, op.getResult().getType(), "result type"))) + } + if (failed(verifyVRegTypeLike(op, op.getResult().getType(), "result type"))) { return failure(); + } if (!allowLowPrecision && failed(verifyNonLowPrecisionVRegElementTypeLike( op.getOperation(), op.getLhs().getType(), "lhs type"))) @@ -5637,27 +6082,33 @@ static LogicalResult verifyTernaryVecOp(TernaryOp op) { } LogicalResult VmaddOp::verify() { - if (failed(verifyTernaryVecOp(*this))) + if (failed(verifyTernaryVecOp(*this))) { return failure(); + } Type elemType = cast(getAcc().getType()).getElementType(); - if (!elemType.isF16() && !elemType.isBF16() && !elemType.isF32()) + if (!elemType.isF16() && !elemType.isBF16() && !elemType.isF32()) { return emitOpError("requires f16/bf16/f32 vector element type"); + } return success(); } LogicalResult VshlOp::verify() { - if (failed(verifyBinaryVecOp(*this))) + if (failed(verifyBinaryVecOp(*this))) { return failure(); + } auto lhsType = cast(getLhs().getType()); - if (!isa(lhsType.getElementType())) + if (!isa(lhsType.getElementType())) { return emitOpError("requires integer vector element type"); + } return success(); } LogicalResult VshrOp::verify() { - if (failed(verifyBinaryVecOp(*this))) + if (failed(verifyBinaryVecOp(*this))) { return failure(); + } auto lhsType = cast(getLhs().getType()); - if (!isa(lhsType.getElementType())) + if (!isa(lhsType.getElementType())) { return emitOpError("requires integer vector element type"); + } return success(); } LogicalResult VaddcOp::verify() { return verifyCarryVecOp(*this); } @@ -5672,8 +6123,9 @@ static LogicalResult verifyReductionVecOp(ReductionOp op) { template static LogicalResult verifyGroupReductionVecOp(ReductionOp op) { - if (failed(verifyReductionVecOp(op))) + if (failed(verifyReductionVecOp(op))) { return failure(); + } auto inputType = cast(op.getInput().getType()); Type elemType = inputType.getElementType(); if (auto intType = dyn_cast(elemType)) { @@ -5683,8 +6135,9 @@ static LogicalResult verifyGroupReductionVecOp(ReductionOp op) { "requires 8-bit, 16-bit, or 32-bit integer vector element type"); return success(); } - if (!elemType.isF16() && !elemType.isF32()) + if (!elemType.isF16() && !elemType.isF32()) { return op.emitOpError("requires i16/i32/f16/f32 vector element type"); + } return success(); } @@ -5692,12 +6145,14 @@ LogicalResult VcgaddOp::verify() { return verifyGroupReductionVecOp(*this); } LogicalResult VcgmaxOp::verify() { return verifyGroupReductionVecOp(*this); } LogicalResult VcgminOp::verify() { return verifyGroupReductionVecOp(*this); } LogicalResult VcpaddOp::verify() { - if (failed(verifyReductionVecOp(*this))) + if (failed(verifyReductionVecOp(*this))) { return failure(); + } auto inputType = cast(getInput().getType()); Type elemType = inputType.getElementType(); - if (!elemType.isF16() && !elemType.isF32()) + if (!elemType.isF16() && !elemType.isF32()) { return emitOpError("requires f16 or f32 vector element type"); + } return success(); } @@ -5720,10 +6175,12 @@ static LogicalResult verifyHistogramOp(HistOp op) { if (!sourceElemType || sourceElemType.getWidth() != 8 || sourceType.getElementCount() != 256) return op.emitOpError("requires source type to be !pto.vreg<256xi8>"); - if (resultType != accType) + if (resultType != accType) { return op.emitOpError("requires result type to match acc type"); - if (!op.getBin().getType().isInteger(32)) + } + if (!op.getBin().getType().isInteger(32)) { return op.emitOpError("requires bin operand to be i32"); + } return success(); } @@ -5746,8 +6203,9 @@ static LogicalResult verifyExtremaPredicateOp(ExtremaOp op) { "requires mask and predicate result to share one mask type"); Type elemType = cast(op.getInput().getType()).getElementType(); - if (elemType.isF16() || elemType.isF32()) + if (elemType.isF16() || elemType.isF32()) { return success(); + } auto intType = dyn_cast(elemType); if (!intType || (intType.getWidth() != 8 && intType.getWidth() != 16 && intType.getWidth() != 32)) @@ -5768,15 +6226,19 @@ static LogicalResult verifyLaneSelectOp(SelectOp op) { auto src0Type = cast(op.getSrc0().getType()); auto src1Type = cast(op.getSrc1().getType()); auto resultType = cast(op.getResult().getType()); - if (src0Type != resultType) + if (src0Type != resultType) { return op.emitOpError("requires src0 and result to have identical vector types"); - if (src1Type.getElementCount() != src0Type.getElementCount()) + } + if (src1Type.getElementCount() != src0Type.getElementCount()) { return op.emitOpError("requires src0/src1 to have identical element counts"); + } auto src1ElemType = dyn_cast(src1Type.getElementType()); - if (!src1ElemType) + if (!src1ElemType) { return op.emitOpError("requires src1 to use integer vector elements"); - if (src1ElemType.getWidth() != getIntOrFloatBitWidth(src0Type.getElementType())) + } + if (src1ElemType.getWidth() != getIntOrFloatBitWidth(src0Type.getElementType())) { return op.emitOpError("requires src1 integer element width to match src0 element width"); + } return success(); } @@ -5803,8 +6265,9 @@ static LogicalResult verifyPartVecOp(PartOp op) { if (op.getLhs().getType() != op.getRhs().getType() || op.getLhs().getType() != op.getResult().getType()) return op.emitOpError("requires operands and result to share one vector type"); - if (!isSupportedPartToken(op.getPart())) + if (!isSupportedPartToken(op.getPart())) { return op.emitOpError("requires part to be LOWER or HIGHER"); + } return success(); } @@ -5833,17 +6296,21 @@ LogicalResult VusqzOp::verify() { failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type")) || failed(verifyVRegTypeLike(*this, getResult().getType(), "result type"))) return failure(); - if (getSrc().getType() != getResult().getType()) + if (getSrc().getType() != getResult().getType()) { return emitOpError("requires src and result to share one vector type"); + } auto srcType = cast(getSrc().getType()); auto elemType = dyn_cast(srcType.getElementType()); - if (!elemType) + if (!elemType) { return emitOpError("requires signed integer vector element type"); - if (elemType.isUnsigned()) + } + if (elemType.isUnsigned()) { return emitOpError("requires signed integer vector element type"); + } unsigned width = elemType.getWidth(); - if (width != 8 && width != 16 && width != 32) + if (width != 8 && width != 16 && width != 32) { return emitOpError("requires s8/s16/s32 vector element type"); + } return success(); } @@ -5851,14 +6318,16 @@ LogicalResult VpackOp::verify() { if (failed(verifyVRegTypeLike(*this, getSrc().getType(), "src type")) || failed(verifyVRegTypeLike(*this, getResult().getType(), "result type"))) return failure(); - if (!isSupportedPartToken(getPart())) + if (!isSupportedPartToken(getPart())) { return emitOpError("requires part to be LOWER or HIGHER"); + } auto srcType = cast(getSrc().getType()); auto resultType = cast(getResult().getType()); Type srcElemType = srcType.getElementType(); Type resultElemType = resultType.getElementType(); - if (!isa(srcElemType) || !isa(resultElemType)) + if (!isa(srcElemType) || !isa(resultElemType)) { return emitOpError("currently requires integer source and result element types"); + } if (resultType.getElementCount() != srcType.getElementCount() * 2) return emitOpError( "requires result element count to be twice the source element count"); @@ -5869,8 +6338,9 @@ LogicalResult VpackOp::verify() { "requires result element width to be half the source element width"); auto srcIntType = cast(srcElemType); auto resultIntType = cast(resultElemType); - if (!resultIntType.isUnsigned()) + if (!resultIntType.isUnsigned()) { return emitOpError("requires unsigned result element type"); + } if (!((srcIntType.getWidth() == 32 && resultIntType.getWidth() == 16) || (srcIntType.getWidth() == 16 && resultIntType.getWidth() == 8))) return emitOpError( @@ -5915,10 +6385,12 @@ LogicalResult VcmpOp::verify() { failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type")) || failed(verifyMaskTypeLike(*this, getResult().getType(), "result type"))) return failure(); - if (getSrc0().getType() != getSrc1().getType()) + if (getSrc0().getType() != getSrc1().getType()) { return emitOpError("requires src0 and src1 to have identical vector types"); - if (!isSupportedCmpMode(getCmpMode())) + } + if (!isSupportedCmpMode(getCmpMode())) { return emitOpError("requires cmp_mode to be one of eq/ne/lt/le/gt/ge"); + } return success(); } @@ -5930,10 +6402,12 @@ LogicalResult VcmpsOp::verify() { auto srcType = cast(getSrc().getType()); Type srcElementType = srcType.getElementType(); Type scalarType = getScalar().getType(); - if (!isCompatibleScalarForSemanticType(srcElementType, scalarType)) + if (!isCompatibleScalarForSemanticType(srcElementType, scalarType)) { return emitOpError("requires scalar type to match source element type"); - if (!isSupportedCmpMode(getCmpMode())) + } + if (!isSupportedCmpMode(getCmpMode())) { return emitOpError("requires cmp_mode to be one of eq/ne/lt/le/gt/ge"); + } return success(); } @@ -5982,25 +6456,31 @@ void VtrcOp::print(OpAsmPrinter &printer) { LogicalResult VtrcOp::verify() { auto inputType = dyn_cast(getInput().getType()); auto resultType = dyn_cast(getResult().getType()); - if (!inputType || !resultType) + if (!inputType || !resultType) { return emitOpError("input and result must be !pto.vreg<...>"); - if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) + } + if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) { return failure(); - if (inputType != resultType) + } + if (inputType != resultType) { return emitOpError("requires input and result to have identical vreg type"); + } auto elemType = inputType.getElementType(); - if (!(elemType.isF16() || elemType.isF32() || elemType.isBF16())) + if (!(elemType.isF16() || elemType.isF32() || elemType.isBF16())) { return emitOpError("requires f16/f32/bf16 vector element type"); + } auto expectedGranularity = getVdupMaskGranularity(elemType); - if (!expectedGranularity) + if (!expectedGranularity) { return emitOpError("requires element type with supported predicate granularity"); + } if (failed(verifyMaskTypeWithGranularityLike(*this, getMask().getType(), "mask type", *expectedGranularity))) return failure(); auto normalized = normalizeRoundModeToken(getRoundMode()); - if (!normalized || !isSupportedVtrcRoundMode(*normalized)) + if (!normalized || !isSupportedVtrcRoundMode(*normalized)) { return emitOpError("round mode must be one of R/A/F/C/Z"); + } return success(); } @@ -6012,14 +6492,17 @@ LogicalResult VmulscvtOp::verify() { auto inputType = cast(getInput().getType()); auto resultType = cast(getResult().getType()); - if (!inputType.getElementType().isF32()) + if (!inputType.getElementType().isF32()) { return emitOpError("requires f32 input vector element type"); - if (!resultType.getElementType().isF16()) + } + if (!resultType.getElementType().isF16()) { return emitOpError("requires f16 result vector element type"); + } auto scalarType = getScalar().getType(); - if (!scalarType.isF32()) + if (!scalarType.isF32()) { return emitOpError("requires f32 scalar operand"); + } if (failed(verifyMaskTypeWithGranularityLike(*this, getMask().getType(), "mask type", "b32"))) @@ -6032,14 +6515,17 @@ LogicalResult VmulscvtOp::verify() { "requires source and result to preserve total vector storage width"); auto normalizedRnd = normalizeRoundModeToken(getRnd()); - if (!normalizedRnd) + if (!normalizedRnd) { return emitOpError("rnd must be one of R/A/F/C/Z/O"); - if (*normalizedRnd != "A") + } + if (*normalizedRnd != "A") { return emitOpError("currently only supports rnd A"); + } auto normalizedPart = normalizeEvenOddPartToken(getPart()); - if (!normalizedPart) + if (!normalizedPart) { return emitOpError("part must be EVEN or ODD"); + } return success(); } @@ -6066,8 +6552,9 @@ ParseResult VcvtOp::parse(OpAsmParser &parser, OperationState &result) { [&](StringRef sourceName, StringRef canonicalName, auto normalizeFn) -> ParseResult { Attribute rawAttr = attrs.get(sourceName); - if (!rawAttr) + if (!rawAttr) { return success(); + } auto strAttr = dyn_cast(rawAttr); if (!strAttr) return parser.emitError(parser.getCurrentLocation()) @@ -6107,28 +6594,33 @@ void VcvtOp::print(OpAsmPrinter &printer) { LogicalResult VcvtOp::verify() { auto inputType = dyn_cast(getInput().getType()); auto resultType = dyn_cast(getResult().getType()); - if (!inputType || !resultType) + if (!inputType || !resultType) { return emitOpError("input and result must be !pto.vreg<...>"); - if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) + } + if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) { return failure(); + } VcvtElemKind inputElemKind = classifyVcvtElemType(inputType.getElementType()); VcvtElemKind resultElemKind = classifyVcvtElemType(resultType.getElementType()); auto contract = lookupVcvtContract(inputElemKind, resultElemKind); - if (!contract) + if (!contract) { return emitOpError("unsupported vcvt source/result element type pair"); + } auto inputElemBits = getVcvtElemBitWidth(inputElemKind); auto resultElemBits = getVcvtElemBitWidth(resultElemKind); - if (!inputElemBits || !resultElemBits) + if (!inputElemBits || !resultElemBits) { return emitOpError("could not determine vcvt element bit width"); + } unsigned maskBitWidth = std::min(*inputElemBits, 32u); StringRef expectedMaskGranularity = maskBitWidth == 8 ? "b8" : maskBitWidth == 16 ? "b16" : maskBitWidth == 32 ? "b32" : ""; - if (expectedMaskGranularity.empty()) + if (expectedMaskGranularity.empty()) { return emitOpError("could not determine vcvt mask granularity"); + } if (failed(verifyMaskTypeWithGranularityLike( *this, getMask().getType(), "mask type", expectedMaskGranularity))) return failure(); @@ -6141,10 +6633,12 @@ LogicalResult VcvtOp::verify() { if (getRndAttr()) { StringRef roundMode = *getRnd(); auto normalizedRoundMode = normalizeRoundModeToken(roundMode); - if (!normalizedRoundMode) + if (!normalizedRoundMode) { return emitOpError("rnd must be one of R/A/F/C/Z/O/H"); - if (!isValidVcvtRoundModeForContract(*normalizedRoundMode, *contract)) + } + if (!isValidVcvtRoundModeForContract(*normalizedRoundMode, *contract)) { return emitOpError("rnd attr is not valid for this vcvt type pair"); + } } if (static_cast(getRndAttr()) != contract->requiresRnd) { return contract->requiresRnd ? emitOpError("requires rnd attr for this vcvt type pair") @@ -6153,8 +6647,9 @@ LogicalResult VcvtOp::verify() { if (getSatAttr()) { StringRef sat = *getSat(); - if (!normalizeSaturationToken(sat)) + if (!normalizeSaturationToken(sat)) { return emitOpError("sat must be SAT or NOSAT"); + } } if (static_cast(getSatAttr()) != contract->requiresSat) { return contract->requiresSat ? emitOpError("requires sat attr for this vcvt type pair") @@ -6164,13 +6659,16 @@ LogicalResult VcvtOp::verify() { if (getPartAttr()) { StringRef part = *getPart(); auto normalizedPart = normalizeVcvtPartToken(part); - if (!normalizedPart) + if (!normalizedPart) { return emitOpError("part must be one of EVEN/ODD/P0/P1/P2/P3"); + } std::optional partFamily = contract->partFamily; - if (!partFamily) + if (!partFamily) { partFamily = classifyVcvtPartFamily(*inputElemBits, *resultElemBits); - if (!partFamily) + } + if (!partFamily) { return emitOpError("part attr is not supported for this vcvt width relation"); + } if (!isValidVcvtPartForFamily(*normalizedPart, *partFamily)) { switch (*partFamily) { case VcvtPartFamily::EvenOdd: @@ -6192,13 +6690,15 @@ LogicalResult VcvtOp::verify() { LogicalResult VbitcastOp::verify() { auto inputType = dyn_cast(getInput().getType()); auto resultType = dyn_cast(getResult().getType()); - if (!inputType || !resultType) + if (!inputType || !resultType) { return emitOpError("input and result must be !pto.vreg<...>"); + } auto getStorageBits = [](VRegType type) -> std::optional { Type elementType = type.getElementType(); - if (auto intType = dyn_cast(elementType)) + if (auto intType = dyn_cast(elementType)) { return type.getElementCount() * static_cast(intType.getWidth()); + } if (auto floatType = dyn_cast(elementType)) return type.getElementCount() * static_cast(floatType.getWidth()); @@ -6207,8 +6707,9 @@ LogicalResult VbitcastOp::verify() { auto inputBits = getStorageBits(inputType); auto resultBits = getStorageBits(resultType); - if (!inputBits || !resultBits) + if (!inputBits || !resultBits) { return emitOpError("requires integer or floating-point vreg element type"); + } if (*inputBits != *resultBits) { return emitOpError("requires source and result vectors to carry the same " "total number of bits"); @@ -6306,10 +6807,12 @@ LogicalResult VmullOp::verify() { return failure(); auto lhsType = cast(getLhs().getType()); auto lhsElemType = dyn_cast(lhsType.getElementType()); - if (!lhsElemType) + if (!lhsElemType) { return emitOpError("requires integer vector element type"); - if (lhsElemType.getWidth() != 32) + } + if (lhsElemType.getWidth() != 32) { return emitOpError("currently requires 32-bit integer vector elements"); + } return success(); } @@ -6341,12 +6844,14 @@ static LogicalResult verifyBinaryVecNoMaskOp(BinaryVecNoMaskOp op) { template static LogicalResult verifyFloatBinaryVecNoMaskOp(BinaryVecNoMaskOp op) { - if (failed(verifyBinaryVecNoMaskOp(op))) + if (failed(verifyBinaryVecNoMaskOp(op))) { return failure(); + } auto lhsType = cast(op.getLhs().getType()); Type elemType = lhsType.getElementType(); - if (!elemType.isF16() && !elemType.isF32()) + if (!elemType.isF16() && !elemType.isF32()) { return op.emitOpError("requires f16 or f32 vector element type"); + } return success(); } @@ -6362,8 +6867,9 @@ static LogicalResult verifyFloatBinaryVecMaskOp(BinaryVecMaskOp op) { return op.emitOpError("requires lhs, rhs, and result to share one vector type"); auto lhsType = cast(op.getLhs().getType()); Type elemType = lhsType.getElementType(); - if (!elemType.isF16() && !elemType.isF32()) + if (!elemType.isF16() && !elemType.isF32()) { return op.emitOpError("requires f16 or f32 vector element type"); + } return success(); } @@ -6378,21 +6884,25 @@ LogicalResult VexpdifOp::verify() { auto inputType = cast(getInput().getType()); auto maxType = cast(getMax().getType()); auto resultType = cast(getResult().getType()); - if (inputType != maxType) + if (inputType != maxType) { return emitOpError("requires input and max to share one vector type"); + } Type inputElemType = inputType.getElementType(); - if (!inputElemType.isF16() && !inputElemType.isF32()) + if (!inputElemType.isF16() && !inputElemType.isF32()) { return emitOpError("requires f16 or f32 input vector element type"); + } auto expectedGranularity = getVdupMaskGranularity(inputElemType); - if (!expectedGranularity) + if (!expectedGranularity) { return emitOpError("requires input element type with supported predicate granularity"); + } if (failed(verifyMaskTypeWithGranularityLike(*this, getMask().getType(), "mask type", *expectedGranularity))) return failure(); - if (!resultType.getElementType().isF32()) + if (!resultType.getElementType().isF32()) { return emitOpError("requires f32 result vector element type"); + } auto inputBits = getVRegStorageBitWidth(inputType); auto resultBits = getVRegStorageBitWidth(resultType); @@ -6401,8 +6911,9 @@ LogicalResult VexpdifOp::verify() { "requires source and result to preserve total vector storage width"); StringRef part = getPart(); - if (part != "EVEN" && part != "ODD") + if (part != "EVEN" && part != "ODD") { return emitOpError("part must be EVEN or ODD"); + } return success(); } @@ -6415,20 +6926,24 @@ LogicalResult VaxpyOp::verify() { auto src0Type = cast(getSrc0().getType()); auto src1Type = cast(getSrc1().getType()); auto resultType = cast(getResult().getType()); - if (src0Type != src1Type || src0Type != resultType) + if (src0Type != src1Type || src0Type != resultType) { return emitOpError("requires src0, src1, and result to share one vector type"); + } Type elemType = src0Type.getElementType(); - if (!elemType.isF16() && !elemType.isF32()) + if (!elemType.isF16() && !elemType.isF32()) { return emitOpError("requires f16 or f32 vector element type"); + } auto expectedGranularity = getVdupMaskGranularity(elemType); - if (!expectedGranularity) + if (!expectedGranularity) { return emitOpError("requires element type with supported predicate granularity"); + } if (failed(verifyMaskTypeWithGranularityLike(*this, getMask().getType(), "mask type", *expectedGranularity))) return failure(); - if (getAlpha().getType() != elemType) + if (getAlpha().getType() != elemType) { return emitOpError("requires alpha type to match vector element type"); + } return success(); } @@ -6441,8 +6956,9 @@ static LogicalResult verifyFusedConvVecOp(ConvOp op) { auto lhsType = cast(op.getLhs().getType()); auto rhsType = cast(op.getRhs().getType()); auto resultType = cast(op.getResult().getType()); - if (lhsType != rhsType) + if (lhsType != rhsType) { return op.emitOpError("requires lhs and rhs to share one vector type"); + } if (!isIntegerOrFloatLike(lhsType.getElementType()) || !isIntegerOrFloatLike(resultType.getElementType())) return op.emitOpError( @@ -6467,19 +6983,24 @@ void Vldsx2Op::getEffects( } LogicalResult Vldsx2Op::verify() { - if (!isBufferLike(getSource().getType())) + if (!isBufferLike(getSource().getType())) { return emitOpError("requires a pointer-like source"); - if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed source"); - if (!getOffset().getType().isIndex()) + } + if (!getOffset().getType().isIndex()) { return emitOpError("requires index offset"); + } if (failed(verifyVRegTypeLike(*this, getLow().getType(), "low result type")) || failed(verifyVRegTypeLike(*this, getHigh().getType(), "high result type"))) return failure(); - if (getLow().getType() != getHigh().getType()) + if (getLow().getType() != getHigh().getType()) { return emitOpError("requires low/high results to share one vector type"); - if (!isSupportedVldx2DistToken(getDist())) + } + if (!isSupportedVldx2DistToken(getDist())) { return emitOpError("requires a supported x2 load distribution token"); + } if (getUpdatedBase() && getUpdatedBase().getType() != getSource().getType()) return emitOpError("requires updated base result to match base type"); @@ -6495,15 +7016,18 @@ void VstsOp::getEffects( template static LogicalResult verifyVstsCommon(StoreOp op) { - if (failed(verifyVRegTypeLike(op, op.getValue().getType(), "value type"))) + if (failed(verifyVRegTypeLike(op, op.getValue().getType(), "value type"))) { return failure(); + } - if (!isBufferLike(op.getDestination().getType())) + if (!isBufferLike(op.getDestination().getType())) { return op.emitOpError("requires a pointer-like destination"); + } MemoryRole destinationRole = classifyMemoryRole(op.getDestination().getType()); - if (destinationRole == MemoryRole::GM) + if (destinationRole == MemoryRole::GM) { return op.emitOpError("requires a UB-backed destination"); + } if (std::optional dist = op.getDist(); dist && !isSupportedVstsDistToken(*dist)) { @@ -6528,8 +7052,9 @@ static LogicalResult verifyVstsCommon(StoreOp op) { } LogicalResult VstsOp::verify() { - if (failed(verifyVstsCommon(*this))) + if (failed(verifyVstsCommon(*this))) { return failure(); + } if (getUpdatedBase() && getUpdatedBase().getType() != getDestination().getType()) return emitOpError("requires updated base result to match base type"); @@ -6548,16 +7073,21 @@ LogicalResult Vstsx2Op::verify() { failed(verifyVRegTypeLike(*this, getHigh().getType(), "high value type")) || failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) return failure(); - if (getLow().getType() != getHigh().getType()) + if (getLow().getType() != getHigh().getType()) { return emitOpError("requires low/high values to share one vector type"); - if (!isBufferLike(getDestination().getType())) + } + if (!isBufferLike(getDestination().getType())) { return emitOpError("requires a pointer-like destination"); - if (classifyMemoryRole(getDestination().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getDestination().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed destination"); - if (!getOffset().getType().isIndex()) + } + if (!getOffset().getType().isIndex()) { return emitOpError("requires index offset"); - if (!isSupportedVstsx2DistToken(getDist())) + } + if (!isSupportedVstsx2DistToken(getDist())) { return emitOpError("requires a supported x2 store distribution token"); + } return success(); } @@ -6569,20 +7099,25 @@ void VscatterOp::getEffects( } LogicalResult VscatterOp::verify() { - if (failed(verifyVRegTypeLike(*this, getValue().getType(), "value type"))) + if (failed(verifyVRegTypeLike(*this, getValue().getType(), "value type"))) { return failure(); - if (!isBufferLike(getDestination().getType())) + } + if (!isBufferLike(getDestination().getType())) { return emitOpError("requires a pointer-like destination"); + } auto offsetsType = dyn_cast(getOffsets().getType()); auto valueType = dyn_cast(getValue().getType()); - if (!offsetsType || !valueType) + if (!offsetsType || !valueType) { return emitOpError("value and offsets must be !pto.vreg<...>"); + } auto offsetsElemType = dyn_cast(offsetsType.getElementType()); - if (!offsetsElemType) + if (!offsetsElemType) { return emitOpError("offset vector must use integer element type"); + } unsigned valueElemWidth = getPTOStorageElemBitWidth(valueType.getElementType()); - if (valueElemWidth != 8 && valueElemWidth != 16 && valueElemWidth != 32) + if (valueElemWidth != 8 && valueElemWidth != 16 && valueElemWidth != 32) { return emitOpError("requires 8-, 16-, or 32-bit value elements"); + } unsigned expectedOffsetWidth = valueElemWidth == 32 ? 32 : 16; if (offsetsElemType.getWidth() != expectedOffsetWidth) @@ -6607,8 +7142,9 @@ LogicalResult VscatterOp::verify() { return emitOpError( "requires destination element type to match value element type"); MemoryRole destinationRole = classifyMemoryRole(getDestination().getType()); - if (destinationRole == MemoryRole::GM) + if (destinationRole == MemoryRole::GM) { return emitOpError("requires a UB-backed destination"); + } return success(); } @@ -6619,17 +7155,21 @@ void VsldbOp::getEffects( } LogicalResult VsldbOp::verify() { - if (!isBufferLike(getSource().getType())) + if (!isBufferLike(getSource().getType())) { return emitOpError("requires a pointer-like source"); - if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getSource().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed source"); + } if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type")) || failed(verifyVRegTypeLike(*this, getResult().getType(), "result type"))) return failure(); - if (!getBlockStride().getType().isSignlessInteger(16)) + if (!getBlockStride().getType().isSignlessInteger(16)) { return emitOpError("requires block_stride to be i16"); - if (!getRepeatStride().getType().isSignlessInteger(16)) + } + if (!getRepeatStride().getType().isSignlessInteger(16)) { return emitOpError("requires repeat_stride to be i16"); + } if (getUpdatedBase() && getUpdatedBase().getType() != getSource().getType()) return emitOpError("requires updated base result to match base type"); @@ -6651,16 +7191,21 @@ void PstiOp::getEffects( } LogicalResult PstiOp::verify() { - if (failed(verifyMaskTypeLike(*this, getValue().getType(), "value type"))) + if (failed(verifyMaskTypeLike(*this, getValue().getType(), "value type"))) { return failure(); - if (!isBufferLike(getDestination().getType())) + } + if (!isBufferLike(getDestination().getType())) { return emitOpError("requires a pointer-like destination"); - if (classifyMemoryRole(getDestination().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getDestination().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed destination"); - if (!matchPattern(getOffset(), m_Constant())) + } + if (!matchPattern(getOffset(), m_Constant())) { return emitOpError("requires offset to be a constant index immediate"); - if (!isSupportedPredicateStoreDist(getDist())) + } + if (!isSupportedPredicateStoreDist(getDist())) { return emitOpError("requires predicate store dist to be NORM or PK"); + } if (getUpdatedBase() && getUpdatedBase().getType() != getDestination().getType()) return emitOpError("requires updated base result to match base type"); @@ -6668,17 +7213,22 @@ LogicalResult PstiOp::verify() { } LogicalResult PstsOp::verify() { - if (failed(verifyMaskTypeLike(*this, getValue().getType(), "value type"))) + if (failed(verifyMaskTypeLike(*this, getValue().getType(), "value type"))) { return failure(); - if (!isBufferLike(getDestination().getType())) + } + if (!isBufferLike(getDestination().getType())) { return emitOpError("requires a pointer-like destination"); + } MemoryRole destinationRole = classifyMemoryRole(getDestination().getType()); - if (destinationRole == MemoryRole::GM) + if (destinationRole == MemoryRole::GM) { return emitOpError("requires a UB-backed destination"); - if (!getOffset().getType().isIndex()) + } + if (!getOffset().getType().isIndex()) { return emitOpError("requires index offset"); - if (!isSupportedPredicateStoreDist(getDist())) + } + if (!isSupportedPredicateStoreDist(getDist())) { return emitOpError("requires predicate store dist to be NORM or PK"); + } if (getUpdatedBase() && getUpdatedBase().getType() != getDestination().getType()) return emitOpError("requires updated base result to match base type"); @@ -6696,14 +7246,18 @@ LogicalResult VsstbOp::verify() { if (failed(verifyVRegTypeLike(*this, getValue().getType(), "value type")) || failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type"))) return failure(); - if (!isBufferLike(getDestination().getType())) + if (!isBufferLike(getDestination().getType())) { return emitOpError("requires a pointer-like destination"); - if (classifyMemoryRole(getDestination().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getDestination().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed destination"); - if (!getBlockStride().getType().isSignlessInteger(16)) + } + if (!getBlockStride().getType().isSignlessInteger(16)) { return emitOpError("requires block_stride to be i16"); - if (!getRepeatStride().getType().isSignlessInteger(16)) + } + if (!getRepeatStride().getType().isSignlessInteger(16)) { return emitOpError("requires repeat_stride to be i16"); + } if (getUpdatedBase() && getUpdatedBase().getType() != getDestination().getType()) return emitOpError("requires updated base result to match base type"); @@ -6718,12 +7272,15 @@ void VstasOp::getEffects( } LogicalResult VstasOp::verify() { - if (failed(verifyStoreAlignChain(getValue(), *this, "value type"))) + if (failed(verifyStoreAlignChain(getValue(), *this, "value type"))) { return failure(); - if (!isBufferLike(getDestination().getType())) + } + if (!isBufferLike(getDestination().getType())) { return emitOpError("requires a pointer-like destination"); - if (classifyMemoryRole(getDestination().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getDestination().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed destination"); + } if (getUpdatedBase() && getUpdatedBase().getType() != getDestination().getType()) return emitOpError("requires updated base result to match base type"); @@ -6738,12 +7295,15 @@ void VstarOp::getEffects( } LogicalResult VstarOp::verify() { - if (failed(verifyStoreAlignChain(getValue(), *this, "value type"))) + if (failed(verifyStoreAlignChain(getValue(), *this, "value type"))) { return failure(); - if (!isBufferLike(getDestination().getType())) + } + if (!isBufferLike(getDestination().getType())) { return emitOpError("requires a pointer-like destination"); - if (classifyMemoryRole(getDestination().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getDestination().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed destination"); + } return success(); } @@ -6760,21 +7320,27 @@ LogicalResult PstuOp::verify() { failed(verifyMaskTypeLike(*this, getValue().getType(), "value type")) || failed(verifyAlignTypeLike(*this, getAlignOut().getType(), "align_out type"))) return failure(); - if (!isBufferLike(getBase().getType()) || !isBufferLike(getBaseOut().getType())) + if (!isBufferLike(getBase().getType()) || !isBufferLike(getBaseOut().getType())) { return emitOpError("requires pointer-like base and base_out"); - if (getBase().getType() != getBaseOut().getType()) + } + if (getBase().getType() != getBaseOut().getType()) { return emitOpError("requires base and base_out to have identical types"); - if (classifyMemoryRole(getBase().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getBase().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed base"); + } auto baseType = cast(getBase().getType()); auto maskType = cast(getValue().getType()); auto elemType = dyn_cast(baseType.getElementType()); - if (!elemType || elemType.isSigned() || (elemType.getWidth() != 16 && elemType.getWidth() != 32)) + if (!elemType || elemType.isSigned() || (elemType.getWidth() != 16 && elemType.getWidth() != 32)) { return emitOpError("requires ui16/ui32 UB base type"); - if (maskType.isB16() && elemType.getWidth() != 16) + } + if (maskType.isB16() && elemType.getWidth() != 16) { return emitOpError("requires !pto.mask to pair with !pto.ptr"); - if (maskType.isB32() && elemType.getWidth() != 32) + } + if (maskType.isB32() && elemType.getWidth() != 32) { return emitOpError("requires !pto.mask to pair with !pto.ptr"); + } return success(); } @@ -6791,12 +7357,15 @@ LogicalResult VstusOp::verify() { failed(verifyVRegTypeLike(*this, getValue().getType(), "value type")) || failed(verifyAlignTypeLike(*this, getAlignOut().getType(), "align_out type"))) return failure(); - if (!isBufferLike(getBase().getType())) + if (!isBufferLike(getBase().getType())) { return emitOpError("requires a pointer-like base"); - if (classifyMemoryRole(getBase().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getBase().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed base"); - if (getBaseOut() && getBaseOut().getType() != getBase().getType()) + } + if (getBaseOut() && getBaseOut().getType() != getBase().getType()) { return emitOpError("requires updated base result to match base type"); + } return success(); } @@ -6813,12 +7382,15 @@ LogicalResult VsturOp::verify() { failed(verifyVRegTypeLike(*this, getValue().getType(), "value type")) || failed(verifyAlignTypeLike(*this, getAlignOut().getType(), "align_out type"))) return failure(); - if (!isBufferLike(getBase().getType())) + if (!isBufferLike(getBase().getType())) { return emitOpError("requires a pointer-like base"); - if (classifyMemoryRole(getBase().getType()) == MemoryRole::GM) + } + if (classifyMemoryRole(getBase().getType()) == MemoryRole::GM) { return emitOpError("requires a UB-backed base"); - if (!isSupportedPostMode(getMode())) + } + if (!isSupportedPostMode(getMode())) { return emitOpError("requires mode to be POST_UPDATE or NO_POST_UPDATE"); + } return success(); } @@ -6839,14 +7411,18 @@ void MteUbGmOp::build(OpBuilder &builder, OperationState &state, Value source, llvm::ArrayRef loops) { state.addOperands({source, destination, lenBurst, nburst.count, nburst.srcStride, nburst.dstStride}); - if (l2CacheCtl) + if (l2CacheCtl) { state.addOperands(l2CacheCtl); - for (const pto::DmaLoopConfig &loop : loops) + } + for (const pto::DmaLoopConfig &loop : loops) { state.addOperands(loop.count); - for (const pto::DmaLoopConfig &loop : loops) + } + for (const pto::DmaLoopConfig &loop : loops) { state.addOperands(loop.srcStride); - for (const pto::DmaLoopConfig &loop : loops) + } + for (const pto::DmaLoopConfig &loop : loops) { state.addOperands(loop.dstStride); + } state.addAttribute( getOperandSegmentSizeAttr(), @@ -6863,10 +7439,12 @@ void MteUbGmOp::build(OpBuilder &builder, OperationState &state, Value source, std::optional loop1, std::optional loop2) { SmallVector loops; - if (loop1) + if (loop1) { loops.push_back(*loop1); - if (loop2) + } + if (loop2) { loops.push_back(*loop2); + } build(builder, state, source, destination, lenBurst, nburst, l2CacheCtl, loops); } @@ -6895,15 +7473,17 @@ ParseResult MteUbGmOp::parse(OpAsmParser &parser, OperationState &result) { if (parseOptionalDmaTripleGroupAlias(parser, {"loop", "loop1", "loop2"}, parsedKeyword, loopGroupOperands)) return failure(); - if (parsedKeyword.empty()) + if (parsedKeyword.empty()) { break; + } loopCountOperands.push_back(loopGroupOperands[0]); loopSrcStrideOperands.push_back(loopGroupOperands[1]); loopDstStrideOperands.push_back(loopGroupOperands[2]); } - if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) + if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) { return failure(); + } Type sourceType, destinationType, lenBurstType, l2CacheCtlType; SmallVector nburstTypes, loopCountTypes, loopSrcStrideTypes, @@ -6914,17 +7494,20 @@ ParseResult MteUbGmOp::parse(OpAsmParser &parser, OperationState &result) { parseDmaTripleTypes(parser, nburstTypes)) return failure(); if (hasL2CacheCtl) { - if (parser.parseComma() || parser.parseType(l2CacheCtlType)) + if (parser.parseComma() || parser.parseType(l2CacheCtlType)) { return failure(); + } } while (succeeded(parser.parseOptionalComma())) { StringRef keyword; - if (parser.parseKeyword(&keyword)) + if (parser.parseKeyword(&keyword)) { return failure(); + } if (isDmaLoopKeyword(keyword)) { SmallVector loopGroupTypes; - if (parseDmaTripleTypes(parser, loopGroupTypes)) + if (parseDmaTripleTypes(parser, loopGroupTypes)) { return failure(); + } loopCountTypes.push_back(loopGroupTypes[0]); loopSrcStrideTypes.push_back(loopGroupTypes[1]); loopDstStrideTypes.push_back(loopGroupTypes[2]); @@ -6977,8 +7560,9 @@ void MteUbGmOp::print(OpAsmPrinter &printer) { << getLenBurst(); printDmaTripleGroup(printer, "nburst", getNBurst(), getNburstSrcStride(), getNburstDstStride()); - if (Value l2CacheCtl = getL2CacheCtl()) + if (Value l2CacheCtl = getL2CacheCtl()) { printer << " l2_cache_ctl(" << l2CacheCtl << ")"; + } for (auto [count, srcStride, dstStride] : llvm::zip(getLoopCounts(), getLoopSrcStrides(), getLoopDstStrides())) printDmaTripleGroup(printer, "loop", count, srcStride, dstStride); @@ -6988,8 +7572,9 @@ void MteUbGmOp::print(OpAsmPrinter &printer) { << ", " << getNburstSrcStride().getType() << ", " << getNburstDstStride().getType(); - if (Value l2CacheCtl = getL2CacheCtl()) + if (Value l2CacheCtl = getL2CacheCtl()) { printer << ", " << l2CacheCtl.getType(); + } for (auto [count, srcStride, dstStride] : llvm::zip(getLoopCounts(), getLoopSrcStrides(), getLoopDstStrides())) printDmaTripleTypes(printer, "loop", count.getType(), srcStride.getType(), @@ -7039,12 +7624,15 @@ void MteGmL1Op::build(OpBuilder &builder, OperationState &state, Value source, state.addOperands( {source, destination, lenBurst, nburst.count, nburst.srcStride, nburst.dstStride}); - for (const pto::DmaLoopConfig &loop : loops) + for (const pto::DmaLoopConfig &loop : loops) { state.addOperands(loop.count); - for (const pto::DmaLoopConfig &loop : loops) + } + for (const pto::DmaLoopConfig &loop : loops) { state.addOperands(loop.srcStride); - for (const pto::DmaLoopConfig &loop : loops) + } + for (const pto::DmaLoopConfig &loop : loops) { state.addOperands(loop.dstStride); + } state.addAttribute( getOperandSegmentSizeAttr(), @@ -7061,10 +7649,12 @@ void MteGmL1Op::build(OpBuilder &builder, OperationState &state, Value source, std::optional loop1, std::optional loop2) { SmallVector loops; - if (loop1) + if (loop1) { loops.push_back(*loop1); - if (loop2) + } + if (loop2) { loops.push_back(*loop2); + } build(builder, state, source, destination, lenBurst, nburst, loops); } @@ -7075,12 +7665,15 @@ void MteL1UbOp::build(OpBuilder &builder, OperationState &state, Value source, state.addOperands( {source, destination, lenBurst, nburst.count, nburst.srcStride, nburst.dstStride}); - for (const pto::DmaLoopConfig &loop : loops) + for (const pto::DmaLoopConfig &loop : loops) { state.addOperands(loop.count); - for (const pto::DmaLoopConfig &loop : loops) + } + for (const pto::DmaLoopConfig &loop : loops) { state.addOperands(loop.srcStride); - for (const pto::DmaLoopConfig &loop : loops) + } + for (const pto::DmaLoopConfig &loop : loops) { state.addOperands(loop.dstStride); + } state.addAttribute( getOperandSegmentSizeAttr(), @@ -7097,10 +7690,12 @@ void MteL1UbOp::build(OpBuilder &builder, OperationState &state, Value source, std::optional loop1, std::optional loop2) { SmallVector loops; - if (loop1) + if (loop1) { loops.push_back(*loop1); - if (loop2) + } + if (loop2) { loops.push_back(*loop2); + } build(builder, state, source, destination, lenBurst, nburst, loops); } @@ -7117,8 +7712,9 @@ void MteGmL1FracOp::build(OpBuilder &builder, OperationState &state, dstGroup.dstLoop3Stride, dstGroup.dstLoop4Stride, ctrl.l2CacheCtrl, ctrl.smallc0En}); bool hasSrcOuterStride = srcLayout.srcOuterStride.has_value(); - if (hasSrcOuterStride) + if (hasSrcOuterStride) { state.addOperands(*srcLayout.srcOuterStride); + } state.addAttribute(getModeAttrName(state.name), CubeLoadFracModeAttr::get(builder.getContext(), mode)); @@ -7141,15 +7737,17 @@ ParseResult MteGmL1Op::parse(OpAsmParser &parser, OperationState &result) { if (parseOptionalDmaTripleGroupAlias(parser, {"loop", "loop1", "loop2"}, parsedKeyword, loopGroupOperands)) return failure(); - if (parsedKeyword.empty()) + if (parsedKeyword.empty()) { break; + } loopCountOperands.push_back(loopGroupOperands[0]); loopSrcStrideOperands.push_back(loopGroupOperands[1]); loopDstStrideOperands.push_back(loopGroupOperands[2]); } - if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) + if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) { return failure(); + } Type sourceType, destinationType, lenBurstType; SmallVector nburstTypes, loopCountTypes, loopSrcStrideTypes, @@ -7161,13 +7759,16 @@ ParseResult MteGmL1Op::parse(OpAsmParser &parser, OperationState &result) { return failure(); while (succeeded(parser.parseOptionalComma())) { StringRef keyword; - if (parser.parseKeyword(&keyword)) + if (parser.parseKeyword(&keyword)) { return failure(); - if (!isDmaLoopKeyword(keyword)) + } + if (!isDmaLoopKeyword(keyword)) { return parser.emitError(parser.getCurrentLocation(), "expected 'loop'"); + } SmallVector loopGroupTypes; - if (parseDmaTripleTypes(parser, loopGroupTypes)) + if (parseDmaTripleTypes(parser, loopGroupTypes)) { return failure(); + } loopCountTypes.push_back(loopGroupTypes[0]); loopSrcStrideTypes.push_back(loopGroupTypes[1]); loopDstStrideTypes.push_back(loopGroupTypes[2]); @@ -7222,15 +7823,17 @@ ParseResult MteL1UbOp::parse(OpAsmParser &parser, OperationState &result) { if (parseOptionalDmaTripleGroupAlias(parser, {"loop", "loop1", "loop2"}, parsedKeyword, loopGroupOperands)) return failure(); - if (parsedKeyword.empty()) + if (parsedKeyword.empty()) { break; + } loopCountOperands.push_back(loopGroupOperands[0]); loopSrcStrideOperands.push_back(loopGroupOperands[1]); loopDstStrideOperands.push_back(loopGroupOperands[2]); } - if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) + if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) { return failure(); + } Type sourceType, destinationType, lenBurstType; SmallVector nburstTypes, loopCountTypes, loopSrcStrideTypes, @@ -7242,13 +7845,16 @@ ParseResult MteL1UbOp::parse(OpAsmParser &parser, OperationState &result) { return failure(); while (succeeded(parser.parseOptionalComma())) { StringRef keyword; - if (parser.parseKeyword(&keyword)) + if (parser.parseKeyword(&keyword)) { return failure(); - if (!isDmaLoopKeyword(keyword)) + } + if (!isDmaLoopKeyword(keyword)) { return parser.emitError(parser.getCurrentLocation(), "expected 'loop'"); + } SmallVector loopGroupTypes; - if (parseDmaTripleTypes(parser, loopGroupTypes)) + if (parseDmaTripleTypes(parser, loopGroupTypes)) { return failure(); + } loopCountTypes.push_back(loopGroupTypes[0]); loopSrcStrideTypes.push_back(loopGroupTypes[1]); loopDstStrideTypes.push_back(loopGroupTypes[2]); @@ -7307,8 +7913,9 @@ ParseResult MteGmL1FracOp::parse(OpAsmParser &parser, OperationState &result) { parseFixedKeywordOperandGroup(parser, "ctrl", 2, ctrlOperands)) return failure(); - if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) + if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) { return failure(); + } Type sourceType, destinationType; SmallVector shapeTypes; @@ -7531,16 +8138,18 @@ void MteGmL1FracOp::print(OpAsmPrinter &printer) { } LogicalResult MteGmL1Op::verify() { - if (failed(verifyCopyGmToUbufOp(*this, true))) + if (failed(verifyCopyGmToUbufOp(*this, true))) { return failure(); + } return verifyDmaLoadStoreLoopGroups( getOperation(), getLoopCounts(), getLoopSrcStrides(), getLoopDstStrides()); } LogicalResult MteL1UbOp::verify() { - if (failed(verifyCopyCbufToUbufLikeOp(*this))) + if (failed(verifyCopyCbufToUbufLikeOp(*this))) { return failure(); + } return verifyDmaLoadStoreLoopGroups( getOperation(), getLoopCounts(), getLoopSrcStrides(), getLoopDstStrides()); @@ -7548,20 +8157,24 @@ LogicalResult MteL1UbOp::verify() { LogicalResult MteL1BtOp::verify() { auto getBufferElementType = [](Type type) -> Type { - if (auto ptrType = dyn_cast(type)) + if (auto ptrType = dyn_cast(type)) { return ptrType.getElementType(); - if (auto memrefType = dyn_cast(type)) + } + if (auto memrefType = dyn_cast(type)) { return memrefType.getElementType(); + } return {}; }; if (!isBufferLike(getSource().getType()) || !isBufferLike(getDestination().getType())) return emitOpError("requires buffer-like source and destination"); - if (getBufferAddressSpace(getSource().getType()) != pto::AddressSpace::MAT) + if (getBufferAddressSpace(getSource().getType()) != pto::AddressSpace::MAT) { return emitOpError("requires MAT source"); - if (getBufferAddressSpace(getDestination().getType()) != pto::AddressSpace::BIAS) + } + if (getBufferAddressSpace(getDestination().getType()) != pto::AddressSpace::BIAS) { return emitOpError("requires BIAS destination"); + } Type srcElem = getBufferElementType(getSource().getType()); Type dstElem = getBufferElementType(getDestination().getType()); @@ -7584,14 +8197,17 @@ LogicalResult MteL1FbOp::verify() { "requires typed !pto.ptr or memref source and destination"); auto getAddressSpace = [](Type type) -> std::optional { - if (auto ptrType = dyn_cast(type)) + if (auto ptrType = dyn_cast(type)) { return ptrType.getMemorySpace().getAddressSpace(); + } if (auto memrefType = dyn_cast(type)) { Attribute memorySpace = memrefType.getMemorySpace(); - if (auto addrSpace = dyn_cast_or_null(memorySpace)) + if (auto addrSpace = dyn_cast_or_null(memorySpace)) { return addrSpace.getAddressSpace(); - if (auto intAttr = dyn_cast_or_null(memorySpace)) + } + if (auto intAttr = dyn_cast_or_null(memorySpace)) { return static_cast(intAttr.getInt()); + } } return std::nullopt; }; @@ -7599,23 +8215,28 @@ LogicalResult MteL1FbOp::verify() { std::optional sourceAS = getAddressSpace(getSource().getType()); std::optional destinationAS = getAddressSpace(getDestination().getType()); - if (!sourceAS || !destinationAS) + if (!sourceAS || !destinationAS) { return emitOpError("requires source and destination with PTO address spaces"); - if (*sourceAS != pto::AddressSpace::MAT) + } + if (*sourceAS != pto::AddressSpace::MAT) { return emitOpError("requires source in mat address space"); - if (*destinationAS != pto::AddressSpace::SCALING) + } + if (*destinationAS != pto::AddressSpace::SCALING) { return emitOpError("requires destination in scaling address space"); + } return success(); } LogicalResult MteGmL1FracOp::verify() { - if (failed(verifyCopyGmToUbufOp(*this, true))) + if (failed(verifyCopyGmToUbufOp(*this, true))) { return failure(); + } auto checkNonNegativeConst = [&](Value value, StringRef name) -> LogicalResult { APInt intValue; - if (matchPattern(value, m_ConstantInt(&intValue)) && intValue.isNegative()) + if (matchPattern(value, m_ConstantInt(&intValue)) && intValue.isNegative()) { return emitOpError() << name << " must be non-negative"; + } return success(); }; if (failed(checkNonNegativeConst(getGroupCount(), "group_count")) || @@ -7779,8 +8400,9 @@ static LogicalResult verifyCubeBridgeLoadStart(Operation *op, Value firstStart, StringRef secondName) { auto checkNonNegativeConst = [&](Value value, StringRef name) -> LogicalResult { APInt intValue; - if (matchPattern(value, m_ConstantInt(&intValue)) && intValue.isNegative()) + if (matchPattern(value, m_ConstantInt(&intValue)) && intValue.isNegative()) { return op->emitOpError() << name << " must be non-negative"; + } return success(); }; @@ -7813,14 +8435,16 @@ struct MxLoadAsmOperand { static std::optional getMxLoadOperandIndex(StringRef keyword, ArrayRef shapeNames) { for (auto [index, name] : llvm::enumerate(shapeNames)) - if (keyword == name) + if (keyword == name) { return index; + } static constexpr StringRef kFullNames[] = { "x_start", "y_start", "x_step", "y_step", "src_stride", "dst_stride"}; for (auto [index, name] : llvm::enumerate(kFullNames)) - if (keyword == name) + if (keyword == name) { return shapeNames.size() + index; + } return std::nullopt; } @@ -7861,20 +8485,24 @@ static ParseResult parseMteL1L0MxOp(OpAsmParser &parser, StringRef keyword; if (succeeded(parser.parseOptionalKeyword(&keyword))) { usesNamedOperands = true; - if (parseNamedOperand(keyword)) + if (parseNamedOperand(keyword)) { return failure(); + } while (succeeded(parser.parseOptionalComma())) { - if (parser.parseKeyword(&keyword) || parseNamedOperand(keyword)) + if (parser.parseKeyword(&keyword) || parseNamedOperand(keyword)) { return failure(); + } } } else { OpAsmParser::UnresolvedOperand operand; - if (parser.parseOperand(operand)) + if (parser.parseOperand(operand)) { return failure(); + } legacyOperands.push_back(operand); while (succeeded(parser.parseOptionalComma())) { - if (parser.parseOperand(operand)) + if (parser.parseOperand(operand)) { return failure(); + } legacyOperands.push_back(operand); } } @@ -7886,8 +8514,9 @@ static ParseResult parseMteL1L0MxOp(OpAsmParser &parser, "expects either four shape-derived or six full " "positional MX operands"); - if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) + if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) { return failure(); + } Type sourceType; Type destinationType; @@ -7899,15 +8528,17 @@ static ParseResult parseMteL1L0MxOp(OpAsmParser &parser, if (usesNamedOperands) { for (unsigned index : namedOperandOrder) { Type type; - if (parser.parseComma() || parser.parseType(type)) + if (parser.parseComma() || parser.parseType(type)) { return failure(); + } namedOperands[index].type = type; } } else { for (size_t index = 0; index < legacyOperands.size(); ++index) { Type type; - if (parser.parseComma() || parser.parseType(type)) + if (parser.parseComma() || parser.parseType(type)) { return failure(); + } legacyTypes.push_back(type); } } @@ -7922,8 +8553,9 @@ static ParseResult parseMteL1L0MxOp(OpAsmParser &parser, if (usesNamedOperands) { for (unsigned index = 0; index < namedOperands.size(); ++index) { - if (!namedOperands[index].present) + if (!namedOperands[index].present) { continue; + } segmentSizes[2 + index] = 1; if (parser.resolveOperand(namedOperands[index].operand, namedOperands[index].type, @@ -7976,14 +8608,16 @@ static void printMteL1L0MxOp(OpAsmPrinter &printer, Operation *operation, } } else { for (auto [index, value] : llvm::enumerate(shapeOperands)) { - if (!value) + if (!value) { continue; + } printer << ", " << shapeNames[index] << "(" << value << ")"; printedOperands.push_back(value); } for (auto [index, value] : llvm::enumerate(fullOperands)) { - if (!value) + if (!value) { continue; + } printer << ", " << fullNames[index] << "(" << value << ")"; printedOperands.push_back(value); } @@ -7992,8 +8626,9 @@ static void printMteL1L0MxOp(OpAsmPrinter &printer, Operation *operation, printer.printOptionalAttrDict(operation->getAttrs(), /*elidedAttrs=*/{"operandSegmentSizes"}); printer << " : " << source.getType() << ", " << destination.getType(); - for (Value value : printedOperands) + for (Value value : printedOperands) { printer << ", " << value.getType(); + } } static LogicalResult verifyMxLoadOperands(Operation *op, @@ -8003,8 +8638,9 @@ static LogicalResult verifyMxLoadOperands(Operation *op, auto checkNonNegativeConst = [&](Value value, StringRef name) -> LogicalResult { APInt intValue; - if (matchPattern(value, m_ConstantInt(&intValue)) && intValue.isNegative()) + if (matchPattern(value, m_ConstantInt(&intValue)) && intValue.isNegative()) { return op->emitOpError() << name << " must be non-negative"; + } return success(); }; @@ -8033,8 +8669,9 @@ static LogicalResult verifyMxLoadOperands(Operation *op, static constexpr StringRef kFullNames[] = { "x_start", "y_start", "x_step", "y_step", "src_stride", "dst_stride"}; for (auto [value, name] : llvm::zip(fullOperands, kFullNames)) - if (!value) + if (!value) { return op->emitOpError() << "full MX form requires " << name; + } if (failed(verifyCubeBridgeLoadStart(op, fullOperands[0], "x_start", fullOperands[1], "y_start"))) return failure(); @@ -8050,13 +8687,15 @@ static LogicalResult verifyMxPointerAlignment(Operation *op, Value pointer, StringRef pointerName, int64_t alignmentBytes) { auto pointerCast = pointer.getDefiningOp(); - if (!pointerCast || !isa(pointerCast.getInput().getType())) + if (!pointerCast || !isa(pointerCast.getInput().getType())) { return success(); + } std::optional address = mlir::getConstantIntValue(pointerCast.getInput()); - if (!address || (*address % alignmentBytes) == 0) + if (!address || (*address % alignmentBytes) == 0) { return success(); + } return op->emitOpError() << "statically known LOAD.MX " << pointerName @@ -8080,8 +8719,9 @@ static LogicalResult verifyExplicitCubeBridgeLoadControls(OpTy op) { auto checkConstRange = [&](Value value, StringRef name, int64_t min, int64_t max) -> LogicalResult { APInt intValue; - if (!matchPattern(value, m_ConstantInt(&intValue))) + if (!matchPattern(value, m_ConstantInt(&intValue))) { return success(); + } int64_t signedValue = intValue.getSExtValue(); if (signedValue < min) return op.emitOpError() @@ -8127,14 +8767,16 @@ LogicalResult MteL0cL1Op::verify() { } LogicalResult MteL1L0aOp::verify() { - if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::LEFT, "LEFT"))) + if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::LEFT, "LEFT"))) { return failure(); + } return verifyCubeBridgeLoadStart(*this); } LogicalResult MteL1L0bOp::verify() { - if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::RIGHT, "RIGHT"))) + if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::RIGHT, "RIGHT"))) { return failure(); + } return verifyCubeBridgeLoadStart(*this); } @@ -8154,8 +8796,9 @@ void MteL1L0aMxOp::print(OpAsmPrinter &printer) { } LogicalResult MteL1L0aMxOp::verify() { - if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::LEFT, "LEFT"))) + if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::LEFT, "LEFT"))) { return failure(); + } if (failed(verifyMxLoadOperands( getOperation(), {getM(), getK(), getStartRow(), getStartCol()}, {"m", "k", "start_row", "start_col"}, @@ -8181,8 +8824,9 @@ void MteL1L0bMxOp::print(OpAsmPrinter &printer) { } LogicalResult MteL1L0bMxOp::verify() { - if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::RIGHT, "RIGHT"))) + if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::RIGHT, "RIGHT"))) { return failure(); + } if (failed(verifyMxLoadOperands( getOperation(), {getK(), getN(), getStartRow(), getStartCol()}, {"k", "n", "start_row", "start_col"}, @@ -8193,30 +8837,34 @@ LogicalResult MteL1L0bMxOp::verify() { } LogicalResult LoadCbufToCaMxOp::verify() { - if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::LEFT, "LEFT"))) + if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::LEFT, "LEFT"))) { return failure(); + } return verifyCubeBridgeLoadStart(getOperation(), getXStartPosition(), "x_start_position", getYStartPosition(), "y_start_position"); } LogicalResult LoadCbufToCbMxOp::verify() { - if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::RIGHT, "RIGHT"))) + if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::RIGHT, "RIGHT"))) { return failure(); + } return verifyCubeBridgeLoadStart(getOperation(), getXStartPosition(), "x_start_position", getYStartPosition(), "y_start_position"); } LogicalResult LoadCbufToCaOp::verify() { - if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::LEFT, "LEFT"))) + if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::LEFT, "LEFT"))) { return failure(); + } return verifyExplicitCubeBridgeLoadControls(*this); } LogicalResult LoadCbufToCbOp::verify() { - if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::RIGHT, "RIGHT"))) + if (failed(verifyCubeBridgeLoadLikeOp(*this, AddressSpace::RIGHT, "RIGHT"))) { return failure(); + } return verifyExplicitCubeBridgeLoadControls(*this); } @@ -8396,18 +9044,21 @@ ParseResult MteL0cUbOp::parse(OpAsmParser &parser, OperationState &result) { parseRequiredOperandWithComma(parser, srcStride) || parseRequiredOperandWithComma(parser, dstStride)) return failure(); - if (parser.parseKeyword("dst_mode") || parser.parseLParen()) + if (parser.parseKeyword("dst_mode") || parser.parseLParen()) { return failure(); + } OptionalParseResult subBlockIdParse = parser.parseOptionalOperand(subBlockId); if (subBlockIdParse.has_value()) { - if (failed(*subBlockIdParse)) + if (failed(*subBlockIdParse)) { return failure(); + } hasSubBlockId = true; } else { StringRef dstModeKeyword; - if (parser.parseKeyword(&dstModeKeyword)) + if (parser.parseKeyword(&dstModeKeyword)) { return failure(); + } if (dstModeKeyword == "split_m") { dstMode = AccStoreUbDstMode::SplitM; } else if (dstModeKeyword == "split_n") { @@ -8419,13 +9070,15 @@ ParseResult MteL0cUbOp::parse(OpAsmParser &parser, OperationState &result) { "dst_mode(split_n)"); } } - if (parser.parseRParen()) + if (parser.parseRParen()) { return failure(); + } if (succeeded(parser.parseOptionalComma()) && parseStructuredAccStoreClauses(parser, state)) return failure(); - if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) + if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) { return failure(); + } Type sourceType, destinationType, mType, nType, srcStrideType, dstStrideType, subBlockIdType; @@ -8438,8 +9091,9 @@ ParseResult MteL0cUbOp::parse(OpAsmParser &parser, OperationState &result) { if (hasSubBlockId && (parser.parseComma() || parser.parseType(subBlockIdType))) return failure(); - if (parseStructuredAccStoreTailTypes(parser, state)) + if (parseStructuredAccStoreTailTypes(parser, state)) { return failure(); + } setStructuredAccStoreSegmentSizes( result, {1, 1, 1, 1, 1, 1, !state.preQuantOperands.empty() ? 1 : 0, @@ -8524,8 +9178,9 @@ void MteL0cUbOp::print(OpAsmPrinter &printer) { printer << " : " << getSource().getType() << ", " << getDestination().getType() << ", " << getM().getType() << ", " << getN().getType() << ", " << getSrcStride().getType() << ", " << getDstStride().getType(); - if (getSubBlockid()) + if (getSubBlockid()) { printer << ", " << getSubBlockid().getType(); + } printStructuredAccStoreOptionalTypes( printer, getPreQuant(), getPreRelu(), getClipValue(), getSplit(), getLoop0SrcStride(), getLoop3Count(), getLoop3SrcStride(), @@ -8552,16 +9207,18 @@ LogicalResult MteL0cUbOp::verify() { return failure(); if (getDstMode() == AccStoreUbDstMode::Single) { - if (!getSubBlockid()) + if (!getSubBlockid()) { return emitOpError("dst_mode(%sub_blockid) requires a sub_blockid operand"); + } APInt subBlockId; if (matchPattern(getSubBlockid(), m_ConstantInt(&subBlockId)) && subBlockId.ugt(1)) return emitOpError("sub_blockid must be 0 or 1"); return success(); } - if (getSubBlockid()) + if (getSubBlockid()) { return emitOpError("split destination modes do not accept sub_blockid"); + } if (getPreQuant() || getPreRelu() || getClipValue() || getPreQuantMode() || getPreReluMode() || getSplit() || getLoop0SrcStride() || @@ -8569,8 +9226,9 @@ LogicalResult MteL0cUbOp::verify() { return emitOpError("dual destination mode cannot be combined with " "pre_quant, pre_relu, clip, nz2dn, nz2nz, or loop3"); } - if (getMode() && *getMode() != AccStoreMode::Nz2nd) + if (getMode() && *getMode() != AccStoreMode::Nz2nd) { return emitOpError("dual destination mode requires normal or nz2nd layout"); + } APInt mValue; APInt nValue; @@ -8810,8 +9468,9 @@ void UBVnotOp::getEffects( } LogicalResult UBVnotOp::verify() { - if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) + if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) { return emitOpError("requires pointer-like operands"); + } if (classifyMemoryRole(getDst().getType()) != MemoryRole::UB || classifyMemoryRole(getSrc().getType()) != MemoryRole::UB) return emitOpError("requires UB-backed operands"); @@ -8830,8 +9489,9 @@ void UBVabsOp::getEffects( } LogicalResult UBVabsOp::verify() { - if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) + if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) { return emitOpError("requires pointer-like operands"); + } if (classifyMemoryRole(getDst().getType()) != MemoryRole::UB || classifyMemoryRole(getSrc().getType()) != MemoryRole::UB) return emitOpError("requires UB-backed operands"); @@ -8850,8 +9510,9 @@ void UBVreluOp::getEffects( } LogicalResult UBVreluOp::verify() { - if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) + if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) { return emitOpError("requires pointer-like operands"); + } if (classifyMemoryRole(getDst().getType()) != MemoryRole::UB || classifyMemoryRole(getSrc().getType()) != MemoryRole::UB) return emitOpError("requires UB-backed operands"); @@ -8891,10 +9552,12 @@ void UBVdupOp::getEffects( } LogicalResult UBVdupOp::verify() { - if (!isBufferLike(getDst().getType())) + if (!isBufferLike(getDst().getType())) { return emitOpError("requires pointer-like dst operand"); - if (classifyMemoryRole(getDst().getType()) != MemoryRole::UB) + } + if (classifyMemoryRole(getDst().getType()) != MemoryRole::UB) { return emitOpError("requires UB-backed dst operand"); + } return success(); } @@ -8910,8 +9573,9 @@ void UBVshlOp::getEffects( } LogicalResult UBVshlOp::verify() { - if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) + if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) { return emitOpError("requires pointer-like operands"); + } if (classifyMemoryRole(getDst().getType()) != MemoryRole::UB || classifyMemoryRole(getSrc().getType()) != MemoryRole::UB) return emitOpError("requires UB-backed operands"); @@ -8930,8 +9594,9 @@ void UBVshrOp::getEffects( } LogicalResult UBVshrOp::verify() { - if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) + if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) { return emitOpError("requires pointer-like operands"); + } if (classifyMemoryRole(getDst().getType()) != MemoryRole::UB || classifyMemoryRole(getSrc().getType()) != MemoryRole::UB) return emitOpError("requires UB-backed operands"); @@ -8950,8 +9615,9 @@ void UBVmulSOp::getEffects( } LogicalResult UBVmulSOp::verify() { - if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) + if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) { return emitOpError("requires pointer-like operands"); + } if (classifyMemoryRole(getDst().getType()) != MemoryRole::UB || classifyMemoryRole(getSrc().getType()) != MemoryRole::UB) return emitOpError("requires UB-backed operands"); diff --git a/lib/PTO/Transforms/BufidSync/BufidSyncAnalysis.h b/lib/PTO/Transforms/BufidSync/BufidSyncAnalysis.h index 6694dbb760..4345758958 100644 --- a/lib/PTO/Transforms/BufidSync/BufidSyncAnalysis.h +++ b/lib/PTO/Transforms/BufidSync/BufidSyncAnalysis.h @@ -179,7 +179,6 @@ inline void printTileGroups(llvm::raw_ostream &os, for (unsigned j = 0; j < tileGroups[i].size(); ++j) { if (j > 0) os << " ; "; printTileValue(os, allTiles[tileGroups[i][j]].tileValue); - //printTileInfo(os, allTiles[tileGroups[i][j]]); } os << "]\n"; } @@ -195,7 +194,6 @@ inline void printVirtualBufIds(llvm::raw_ostream &os, for (unsigned i = 0; i < vbid.tiles.size(); ++i) { if (i > 0) os << " ; "; printTileValue(os, vbid.tiles[i].tileValue); - //printTileInfo(os, vbid.tiles[i]); } os << "]\n"; } diff --git a/lib/PTO/Transforms/BufidSync/BufidSyncIdAlloc.cpp b/lib/PTO/Transforms/BufidSync/BufidSyncIdAlloc.cpp index bbdb73a1de..9c3d8cd7bd 100644 --- a/lib/PTO/Transforms/BufidSync/BufidSyncIdAlloc.cpp +++ b/lib/PTO/Transforms/BufidSync/BufidSyncIdAlloc.cpp @@ -308,7 +308,7 @@ void BufidSyncIdAlloc::reuseIds() { }; int iteration = 0; - while (maxPhysicalIdUsed_ >= (int)physicalBufIdCount_) { + while (maxPhysicalIdUsed_ >= static_cast(physicalBufIdCount_)) { ++iteration; DenseMap> logicIdPipes; diff --git a/lib/PTO/Transforms/BufidSync/BufidSyncIdAlloc.h b/lib/PTO/Transforms/BufidSync/BufidSyncIdAlloc.h index ebaa030682..0dda2471cd 100644 --- a/lib/PTO/Transforms/BufidSync/BufidSyncIdAlloc.h +++ b/lib/PTO/Transforms/BufidSync/BufidSyncIdAlloc.h @@ -27,7 +27,7 @@ class BufidSyncIdAlloc { void computeLifeIntervals(); void linearScanAllocate(); - bool needsReuse() const { return maxPhysicalIdUsed_ >= (int)physicalBufIdCount_; } + bool needsReuse() const { return maxPhysicalIdUsed_ >= static_cast(physicalBufIdCount_); } void reuseIds(); void compactPhysicalIds(); bool validateNoSamePhysicalIdNesting(std::string *error = nullptr) const; diff --git a/lib/PTO/Transforms/ConvertToPTOOp.cpp b/lib/PTO/Transforms/ConvertToPTOOp.cpp index 9fb56eef67..dc273ac4ab 100644 --- a/lib/PTO/Transforms/ConvertToPTOOp.cpp +++ b/lib/PTO/Transforms/ConvertToPTOOp.cpp @@ -39,15 +39,17 @@ namespace { //===---------------------------------------------------------------------===// std::optional getPadValue(std::optional maybeAlloc) { - if (!maybeAlloc.has_value()) + if (!maybeAlloc.has_value()) { return std::nullopt; + } return std::nullopt; } std::optional getLeftPadNum(PatternRewriter &rewriter, std::optional maybeAlloc) { - if (!maybeAlloc.has_value()) + if (!maybeAlloc.has_value()) { return std::nullopt; + } for (auto *user : maybeAlloc.value()->getUsers()) { if (auto subviewOp = llvm::dyn_cast(user)) { @@ -75,14 +77,16 @@ std::pair, std::optional> getUniqueInitInfo(PatternRewriter &rewriter, std::optional maybeAlloc, pto::TLoadOp loadOp) { - if (!maybeAlloc.has_value()) + if (!maybeAlloc.has_value()) { return {std::nullopt, std::nullopt}; + } std::optional initOp = std::nullopt; std::optional initCondition = std::nullopt; for (auto *user : (*maybeAlloc)->getUsers()) { - if (llvm::isa(user)) + if (llvm::isa(user)) { continue; + } auto maybeInitOp = getInitInfo(user, loadOp).first; if (maybeInitOp.has_value() && !initOp.has_value()) { std::tie(initOp, initCondition) = getInitInfo(user, loadOp); diff --git a/lib/PTO/Transforms/CppPostprocess.cpp b/lib/PTO/Transforms/CppPostprocess.cpp index 554193a744..e04bce8edb 100644 --- a/lib/PTO/Transforms/CppPostprocess.cpp +++ b/lib/PTO/Transforms/CppPostprocess.cpp @@ -30,8 +30,9 @@ struct ParsedMarkerCall { static bool parseMarkerArgs(llvm::StringRef argsRef, llvm::SmallVectorImpl &args) { args.clear(); - if (argsRef.empty()) + if (argsRef.empty()) { return true; + } int parenDepth = 0; size_t partBegin = 0; @@ -42,8 +43,9 @@ static bool parseMarkerArgs(llvm::StringRef argsRef, continue; } if (c == ')') { - if (parenDepth > 0) + if (parenDepth > 0) { --parenDepth; + } continue; } if (c == ',' && parenDepth == 0) { @@ -51,8 +53,9 @@ static bool parseMarkerArgs(llvm::StringRef argsRef, partBegin = i + 1; } } - if (partBegin > argsRef.size()) + if (partBegin > argsRef.size()) { return false; + } args.push_back(argsRef.drop_front(partBegin).trim()); return true; } @@ -61,18 +64,21 @@ static bool parseLastUseMarkerName(llvm::StringRef markerName, std::string &callee, std::string &lastUseArgs) { static constexpr llvm::StringLiteral kPrefix = "PTOAS__LAST_USE__"; - if (!markerName.starts_with(kPrefix)) + if (!markerName.starts_with(kPrefix)) { return false; + } llvm::StringRef payload = markerName.drop_front(kPrefix.size()); size_t split = payload.find("__"); - if (split == llvm::StringRef::npos) + if (split == llvm::StringRef::npos) { return false; + } callee = payload.take_front(split).str(); llvm::StringRef encoded = payload.drop_front(split + 2); - if (callee.empty() || encoded.empty()) + if (callee.empty() || encoded.empty()) { return false; + } lastUseArgs.clear(); size_t pos = 0; @@ -81,8 +87,9 @@ static bool parseLastUseMarkerName(llvm::StringRef markerName, llvm::StringRef token = next == llvm::StringRef::npos ? encoded.drop_front(pos) : encoded.slice(pos, next); - if (token.empty()) + if (token.empty()) { return false; + } if (!llvm::all_of(token, [](char c) { return std::isdigit(c); })) return false; if (!lastUseArgs.empty()) @@ -103,8 +110,9 @@ bool rewriteLastUseMarkersInCpp(std::string &cpp) { static constexpr llvm::StringLiteral kPrefix = "PTOAS__LAST_USE__"; while (true) { size_t markerPos = cpp.find(kPrefix.str(), searchPos); - if (markerPos == std::string::npos) + if (markerPos == std::string::npos) { break; + } size_t lparenPos = markerPos + kPrefix.size(); while (lparenPos < cpp.size() && cpp[lparenPos] != '(') @@ -159,8 +167,9 @@ bool rewriteLastUseMarkersInCpp(std::string &cpp) { replacement.append(callee); replacement.push_back('('); for (size_t i = 0; i < call.args.size(); ++i) { - if (i) + if (i) { replacement.append(", "); + } replacement.append(call.args[i].str()); } replacement.push_back(')'); diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index f160a275e4..7fb7d9a386 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -120,8 +120,9 @@ struct OperandTypeInfo { /// Equality for SpecKey caching — only compares fields relevant to each kind. bool operator==(const OperandTypeInfo &rhs) const { - if (kind != rhs.kind || dtype != rhs.dtype) + if (kind != rhs.kind || dtype != rhs.dtype) { return false; + } if (kind == OperandKind::Tile) return tileShape == rhs.tileShape && tileValidShape == rhs.tileValidShape && @@ -129,10 +130,12 @@ struct OperandTypeInfo { blayout == rhs.blayout && slayout == rhs.slayout && fractal == rhs.fractal && pad == rhs.pad && compact == rhs.compact; - if (kind == OperandKind::Vector) + if (kind == OperandKind::Vector) { return vectorShape == rhs.vectorShape; - if (kind == OperandKind::Scalar) + } + if (kind == OperandKind::Scalar) { return scalarValue == rhs.scalarValue; + } return viewShape == rhs.viewShape && viewStrides == rhs.viewStrides && viewMemorySpace == rhs.viewMemorySpace && @@ -167,31 +170,39 @@ struct SpecKeyInfo : public llvm::DenseMapInfo { if (op.kind == OperandKind::Tile) { h = llvm::hash_combine(h, op.tileMemorySpace, op.blayout, op.slayout, op.fractal, op.pad, op.compact); - for (int64_t d : op.tileShape) + for (int64_t d : op.tileShape) { h = llvm::hash_combine(h, d); - for (int64_t d : op.tileValidShape) + } + for (int64_t d : op.tileValidShape) { h = llvm::hash_combine(h, d); + } } else if (op.kind == OperandKind::Vector) { - for (int64_t d : op.vectorShape) + for (int64_t d : op.vectorShape) { h = llvm::hash_combine(h, d); + } } else if (op.kind == OperandKind::Scalar) { h = llvm::hash_combine(h, op.scalarValue.has_value()); - if (op.scalarValue) + if (op.scalarValue) { h = llvm::hash_combine(h, *op.scalarValue); + } } if (op.kind == OperandKind::View) { h = llvm::hash_combine(h, op.viewMemorySpace); - for (int64_t d : op.viewShape) + for (int64_t d : op.viewShape) { h = llvm::hash_combine(h, d); - for (int64_t d : op.viewStrides) + } + for (int64_t d : op.viewStrides) { h = llvm::hash_combine(h, d); + } h = llvm::hash_combine(h, op.viewLayout.has_value()); - if (op.viewLayout) + if (op.viewLayout) { h = llvm::hash_combine(h, static_cast(*op.viewLayout)); + } } } - for (const auto &[attrName, attrValue] : key.contextAttrs) + for (const auto &[attrName, attrValue] : key.contextAttrs) { h = llvm::hash_combine(h, attrName, attrValue); + } return h; } static bool isEqual(const SpecKey &lhs, const SpecKey &rhs) { @@ -202,28 +213,72 @@ struct SpecKeyInfo : public llvm::DenseMapInfo { // Helpers // ============================================================================ static std::string getDtypeString(Type elemTy) { - if (elemTy.isIndex()) return "i32"; - if (elemTy.isInteger(1)) return "i1"; - if (elemTy.isF32()) return "f32"; - if (elemTy.isF16()) return "f16"; - if (elemTy.isBF16()) return "bf16"; - if (isa(elemTy)) return "f8e4m3"; - if (isa(elemTy)) return "f8e5m2"; - if (isa(elemTy)) return "hif8"; - if (isa(elemTy)) return "f4e1m2x2"; - if (isa(elemTy)) return "f4e2m1x2"; - if (elemTy.isUnsignedInteger(64)) return "ui64"; - if (elemTy.isUnsignedInteger(32)) return "ui32"; - if (elemTy.isUnsignedInteger(16)) return "ui16"; - if (elemTy.isUnsignedInteger(8)) return "ui8"; - if (elemTy.isSignedInteger(64)) return "si64"; - if (elemTy.isSignedInteger(32)) return "si32"; - if (elemTy.isSignedInteger(16)) return "si16"; - if (elemTy.isSignedInteger(8)) return "si8"; - if (elemTy.isSignlessInteger(64)) return "i64"; - if (elemTy.isSignlessInteger(32)) return "i32"; - if (elemTy.isSignlessInteger(16)) return "i16"; - if (elemTy.isSignlessInteger(8)) return "i8"; + if (elemTy.isIndex()) { + return "i32"; + } + if (elemTy.isInteger(1)) { + return "i1"; + } + if (elemTy.isF32()) { + return "f32"; + } + if (elemTy.isF16()) { + return "f16"; + } + if (elemTy.isBF16()) { + return "bf16"; + } + if (isa(elemTy)) { + return "f8e4m3"; + } + if (isa(elemTy)) { + return "f8e5m2"; + } + if (isa(elemTy)) { + return "hif8"; + } + if (isa(elemTy)) { + return "f4e1m2x2"; + } + if (isa(elemTy)) { + return "f4e2m1x2"; + } + if (elemTy.isUnsignedInteger(64)) { + return "ui64"; + } + if (elemTy.isUnsignedInteger(32)) { + return "ui32"; + } + if (elemTy.isUnsignedInteger(16)) { + return "ui16"; + } + if (elemTy.isUnsignedInteger(8)) { + return "ui8"; + } + if (elemTy.isSignedInteger(64)) { + return "si64"; + } + if (elemTy.isSignedInteger(32)) { + return "si32"; + } + if (elemTy.isSignedInteger(16)) { + return "si16"; + } + if (elemTy.isSignedInteger(8)) { + return "si8"; + } + if (elemTy.isSignlessInteger(64)) { + return "i64"; + } + if (elemTy.isSignlessInteger(32)) { + return "i32"; + } + if (elemTy.isSignlessInteger(16)) { + return "i16"; + } + if (elemTy.isSignlessInteger(8)) { + return "i8"; + } return ""; } @@ -232,10 +287,12 @@ static std::string getDtypeString(Type elemTy) { static Value bridgeOperandToType(OpBuilder &builder, Location loc, Value operand, Type dstTy) { Type srcTy = operand.getType(); - if (srcTy == dstTy) + if (srcTy == dstTy) { return operand; - if (srcTy.isIndex() && isa(dstTy)) + } + if (srcTy.isIndex() && isa(dstTy)) { return builder.create(loc, dstTy, operand); + } return builder.create(loc, dstTy, operand) .getResult(0); } @@ -245,12 +302,14 @@ static StringRef getTileOpName(Operation *op) { } static std::string getTargetArchString(Operation *op) { - if (!op) + if (!op) { return ""; + } for (ModuleOp current = op->getParentOfType(); current; current = current->getParentOfType()) { - if (auto targetAttr = current->getAttrOfType("pto.target_arch")) + if (auto targetAttr = current->getAttrOfType("pto.target_arch")) { return targetAttr.getValue().str(); + } } return ""; } @@ -289,37 +348,44 @@ static std::string getMemorySpaceString(MemRefType mrTy) { } static std::string getBLayoutString(int32_t blayout) { - if (blayout == static_cast(pto::BLayout::ColMajor)) + if (blayout == static_cast(pto::BLayout::ColMajor)) { return "col_major"; + } return "row_major"; } static std::string getSLayoutString(int32_t slayout) { - if (slayout == static_cast(pto::SLayout::RowMajor)) + if (slayout == static_cast(pto::SLayout::RowMajor)) { return "row_major"; - if (slayout == static_cast(pto::SLayout::ColMajor)) + } + if (slayout == static_cast(pto::SLayout::ColMajor)) { return "col_major"; + } return "none_box"; } static constexpr llvm::StringLiteral kLayoutAttrName = "layout"; static std::optional getLayoutAttrFromOp(Operation *op) { - if (!op) + if (!op) { return std::nullopt; - if (auto attr = op->getAttrOfType(kLayoutAttrName)) + } + if (auto attr = op->getAttrOfType(kLayoutAttrName)) { return attr.getLayout(); + } return std::nullopt; } static std::optional resolveViewLayout(Value value) { - if (!value) + if (!value) { return std::nullopt; + } Operation *def = value.getDefiningOp(); while (def) { - if (auto layout = getLayoutAttrFromOp(def)) + if (auto layout = getLayoutAttrFromOp(def)) { return layout; + } if (auto subview = dyn_cast(def)) { value = subview.getSource(); def = value.getDefiningOp(); @@ -346,8 +412,9 @@ static std::optional resolveViewLayout(Value value) { } static std::optional getViewLayoutString(std::optional layout) { - if (!layout) + if (!layout) { return std::nullopt; + } return stringifyLayout(*layout).str(); } @@ -454,8 +521,9 @@ static bool tryAppendPrecisionType( SmallVectorImpl> &attrs, PrecisionT highPrecision) { auto typed = dyn_cast(op); - if (!typed) + if (!typed) { return false; + } PrecisionT precision = typed.getPrecisionType(); attrs.emplace_back("precisionType", getPrecisionTypeString(precision).str()); @@ -479,11 +547,13 @@ static LogicalResult appendOpContextAttrs( SmallVectorImpl> &attrs) { if (auto tcvt = dyn_cast(op)) { std::optional roundMode = getTCvtRoundModeString(tcvt); - if (roundMode) + if (roundMode) { attrs.emplace_back("round_mode", *roundMode); + } } - if (auto trandom = dyn_cast(op)) + if (auto trandom = dyn_cast(op)) { attrs.emplace_back("rounds", getTRandomRoundsString(trandom)); + } if (auto tcmp = dyn_cast(op)) { if (auto cmpModeAttr = tcmp.getCmpModeAttr()) { attrs.emplace_back("cmp_mode", @@ -519,8 +589,9 @@ static LogicalResult appendOpContextAttrs( } if (auto thistogram = dyn_cast(op)) { int byte = 1; - if (auto byteAttr = thistogram.getByteAttr()) + if (auto byteAttr = thistogram.getByteAttr()) { byte = byteAttr.getInt(); + } attrs.emplace_back("byte", std::to_string(byte)); } if (auto tci = dyn_cast(op)) { @@ -593,14 +664,16 @@ static bool getStaticIntFromValue(Value value, int64_t &out) { static int64_t getStaticIntOrDynamic(OpFoldResult ofr) { if (isa(ofr)) { Attribute attr = cast(ofr); - if (auto intAttr = dyn_cast(attr)) + if (auto intAttr = dyn_cast(attr)) { return intAttr.getInt(); + } return ShapedType::kDynamic; } Value value = cast(ofr); int64_t result = ShapedType::kDynamic; - if (getStaticIntFromValue(value, result)) + if (getStaticIntFromValue(value, result)) { return result; + } return ShapedType::kDynamic; } @@ -608,8 +681,9 @@ static void recordStaticSizes(ArrayRef inputs, SmallVectorImpl &out) { out.clear(); out.reserve(inputs.size()); - for (OpFoldResult ofr : inputs) + for (OpFoldResult ofr : inputs) { out.push_back(getStaticIntOrDynamic(ofr)); + } } static SmallVector combineSubviewStrides(ArrayRef baseStrides, @@ -631,8 +705,9 @@ static SmallVector combineSubviewStrides(ArrayRef baseStrides, static void populateViewShapeAndStrides(Value value, SmallVectorImpl &shape, SmallVectorImpl &strides) { - if (!value) + if (!value) { return; + } if (auto partition = value.getDefiningOp()) { populateViewShapeAndStrides(partition.getSource(), shape, strides); @@ -667,10 +742,12 @@ static void populateViewShapeAndStrides(Value value, populateViewShapeAndStrides(subview.getSource(), shape, strides); SmallVector subviewShape; recordStaticSizes(subview.getMixedSizes(), subviewShape); - if (!subviewShape.empty()) + if (!subviewShape.empty()) { shape = subviewShape; - if (!strides.empty()) + } + if (!strides.empty()) { strides = combineSubviewStrides(strides, subview.getMixedStrides()); + } return; } @@ -678,11 +755,13 @@ static void populateViewShapeAndStrides(Value value, if (shape.empty()) { SmallVector reinterpretShape; recordStaticSizes(reinterpret.getMixedSizes(), reinterpretShape); - if (!reinterpretShape.empty()) + if (!reinterpretShape.empty()) { shape = reinterpretShape; + } } - if (strides.empty()) + if (strides.empty()) { recordStaticSizes(reinterpret.getMixedStrides(), strides); + } return; } @@ -692,8 +771,9 @@ static void populateViewShapeAndStrides(Value value, } if (auto memrefTy = dyn_cast(value.getType())) { - if (shape.empty()) + if (shape.empty()) { shape.assign(memrefTy.getShape().begin(), memrefTy.getShape().end()); + } if (strides.empty()) { int64_t offset = ShapedType::kDynamic; if (succeeded( @@ -711,12 +791,14 @@ static std::optional buildOperandTypeInfo(Value value) { OperandTypeInfo info; info.kind = OperandKind::Tile; info.dtype = getDtypeString(tbTy.getElementType()); - if (info.dtype.empty()) + if (info.dtype.empty()) { return std::nullopt; + } info.tileShape.assign(tbTy.getShape().begin(), tbTy.getShape().end()); auto validShape = tbTy.getValidShape(); - if (validShape.empty()) + if (validShape.empty()) { info.tileValidShape.assign(tbTy.getShape().begin(), tbTy.getShape().end()); + } else info.tileValidShape.assign(validShape.begin(), validShape.end()); info.tileMemorySpace = getMemorySpaceString(tbTy); @@ -738,13 +820,15 @@ static std::optional buildOperandTypeInfo(Value value) { OperandTypeInfo info; info.kind = OperandKind::View; info.dtype = getDtypeString(mrTy.getElementType()); - if (info.dtype.empty()) + if (info.dtype.empty()) { return std::nullopt; + } info.viewMemorySpace = getMemorySpaceString(mrTy); info.viewLayout = resolveViewLayout(value); populateViewShapeAndStrides(value, info.viewShape, info.viewStrides); - if (info.viewShape.empty()) + if (info.viewShape.empty()) { info.viewShape.assign(mrTy.getShape().begin(), mrTy.getShape().end()); + } if (info.viewStrides.empty()) { int64_t offset = ShapedType::kDynamic; if (succeeded(mlir::pto::getPTOMemRefStridesAndOffset( @@ -759,15 +843,18 @@ static std::optional buildOperandTypeInfo(Value value) { OperandTypeInfo info; info.kind = OperandKind::View; info.dtype = getDtypeString(viewTy.getElementType()); - if (info.dtype.empty()) + if (info.dtype.empty()) { return std::nullopt; + } info.viewMemorySpace = "gm"; info.viewLayout = resolveViewLayout(value); populateViewShapeAndStrides(value, info.viewShape, info.viewStrides); - if (info.viewShape.empty()) + if (info.viewShape.empty()) { info.viewShape.assign(viewTy.getShape().begin(), viewTy.getShape().end()); - if (info.viewStrides.empty()) + } + if (info.viewStrides.empty()) { info.viewStrides.assign(viewTy.getRank(), ShapedType::kDynamic); + } return info; } @@ -776,8 +863,9 @@ static std::optional buildOperandTypeInfo(Value value) { OperandTypeInfo info; info.kind = OperandKind::Vector; info.dtype = getDtypeString(vecTy.getElementType()); - if (info.dtype.empty()) + if (info.dtype.empty()) { return std::nullopt; + } info.vectorShape.assign(vecTy.getShape().begin(), vecTy.getShape().end()); return info; } @@ -786,11 +874,13 @@ static std::optional buildOperandTypeInfo(Value value) { OperandTypeInfo info; info.kind = OperandKind::Scalar; info.dtype = getDtypeString(ty); - if (info.dtype.empty()) + if (info.dtype.empty()) { return std::nullopt; + } int64_t scalarValue = 0; - if (getStaticIntFromValue(value, scalarValue)) + if (getStaticIntFromValue(value, scalarValue)) { info.scalarValue = scalarValue; + } return info; } @@ -814,8 +904,9 @@ static FailureOr buildSpecKey(Operation *op) { return failure(); } - if (failed(appendOpContextAttrs(op, key.contextAttrs))) + if (failed(appendOpContextAttrs(op, key.contextAttrs))) { return failure(); + } return key; } @@ -850,8 +941,9 @@ struct ExpandTileOpPass static void appendJsonIntArray(std::string &json, ArrayRef arr) { json += "["; for (size_t i = 0; i < arr.size(); ++i) { - if (i > 0) + if (i > 0) { json += ","; + } json += std::to_string(arr[i]); } json += "]"; @@ -862,8 +954,9 @@ static void appendJsonDimArray(std::string &json, ArrayRef arr, bool negativeIsDynamic = false) { json += "["; for (size_t i = 0; i < arr.size(); ++i) { - if (i > 0) + if (i > 0) { json += ","; + } int64_t dim = arr[i]; if (ShapedType::isDynamic(dim) || (negativeIsDynamic && dim < 0)) { json += "null"; @@ -878,8 +971,9 @@ static std::string buildOperandSpecsJson(const SpecKey &key) { std::string json = "["; for (size_t i = 0; i < key.operands.size(); ++i) { const auto &op = key.operands[i]; - if (i > 0) + if (i > 0) { json += ","; + } if (op.kind == OperandKind::Tile) { json += "{\"kind\":\"tile\",\"dtype\":\"" + op.dtype + "\",\"shape\":"; @@ -909,10 +1003,12 @@ static std::string buildOperandSpecsJson(const SpecKey &key) { if (!op.viewStrides.empty()) { json += ",\"strides\":["; for (size_t dim = 0; dim < op.viewStrides.size(); ++dim) { - if (dim > 0) + if (dim > 0) { json += ","; - if (ShapedType::isDynamic(op.viewStrides[dim])) + } + if (ShapedType::isDynamic(op.viewStrides[dim])) { json += "null"; + } else json += std::to_string(op.viewStrides[dim]); } @@ -948,8 +1044,9 @@ static std::string buildOperandSpecsJson(const SpecKey &key) { } static std::string dimSuffix(int64_t dim) { - if (ShapedType::isDynamic(dim)) + if (ShapedType::isDynamic(dim)) { return "d"; + } return std::to_string(dim); } @@ -962,10 +1059,12 @@ static std::string buildUniqueFunctionBaseName(const SpecKey &key) { : "_scalar"; uniqueName += "_" + op.dtype; if (op.kind == OperandKind::Tile) { - for (int64_t d : op.tileShape) + for (int64_t d : op.tileShape) { uniqueName += "_" + std::to_string(d); - for (int64_t d : op.tileValidShape) + } + for (int64_t d : op.tileValidShape) { uniqueName += "_v" + std::to_string(d); + } uniqueName += "_bl" + std::to_string(op.blayout); uniqueName += "_sl" + std::to_string(op.slayout); uniqueName += "_fr" + std::to_string(op.fractal); @@ -974,30 +1073,36 @@ static std::string buildUniqueFunctionBaseName(const SpecKey &key) { } else if (op.kind == OperandKind::View) { uniqueName += "_ms_" + op.viewMemorySpace; uniqueName += "_shape"; - for (int64_t d : op.viewShape) + for (int64_t d : op.viewShape) { uniqueName += "_" + dimSuffix(d); + } uniqueName += "_strides"; - for (int64_t d : op.viewStrides) + for (int64_t d : op.viewStrides) { uniqueName += "_" + dimSuffix(d); - if (op.viewLayout) + } + if (op.viewLayout) { uniqueName += "_vl_" + stringifyLayout(*op.viewLayout).str(); + } } else if (op.kind == OperandKind::Vector) { - for (int64_t d : op.vectorShape) + for (int64_t d : op.vectorShape) { uniqueName += "_" + std::to_string(d); + } } else if (op.kind == OperandKind::Scalar && op.scalarValue) { uniqueName += "_sv" + std::to_string(*op.scalarValue); } } - for (const auto &[attrName, attrValue] : key.contextAttrs) + for (const auto &[attrName, attrValue] : key.contextAttrs) { uniqueName += "_ctx_" + attrName + "_" + attrValue; + } return uniqueName; } static std::string buildUniqueFunctionName(const SpecKey &key, StringRef candidateId) { std::string uniqueName = buildUniqueFunctionBaseName(key); - if (!candidateId.empty()) + if (!candidateId.empty()) { uniqueName += "__" + candidateId.str(); + } return uniqueName; } @@ -1005,8 +1110,9 @@ static std::string buildContextAttrsJson(const SpecKey &key) { std::string json = "{"; for (size_t i = 0; i < key.contextAttrs.size(); ++i) { const auto &[attrName, attrValue] = key.contextAttrs[i]; - if (i > 0) + if (i > 0) { json += ","; + } json += "\""; json += attrName; json += "\":\""; @@ -1027,8 +1133,9 @@ func::FuncOp ExpandState::invokeInProcessTileLib(const SpecKey &key, const std::string &uniqueName, ModuleOp mod, MLIRContext *ctx) { - if (!tileLibService) + if (!tileLibService) { return nullptr; + } pto::TileLibMaterializationRequest request; request.target = key.targetArch; @@ -1054,8 +1161,9 @@ func::FuncOp ExpandState::invokeInProcessTileLib(const SpecKey &key, } SmallVector sourceFuncs; - for (func::FuncOp fn : sourceModule.getOps()) + for (func::FuncOp fn : sourceModule.getOps()) { sourceFuncs.push_back(fn); + } if (sourceFuncs.empty()) { llvm::errs() << "ExpandTileOp: in-process PTODSL returned no func.func\n"; return failure(); @@ -1094,8 +1202,9 @@ func::FuncOp ExpandState::invokeInProcessTileLib(const SpecKey &key, llvm::errs() << "ExpandTileOp: failed to rewrite imported symbol @" << renamed.getKey() << " in @" << fn.getSymName() << "\n"; - for (func::FuncOp imported : clonedFuncs) + for (func::FuncOp imported : clonedFuncs) { imported.erase(); + } return failure(); } } @@ -1151,8 +1260,9 @@ func::FuncOp ExpandState::invokeTileLib(const SpecKey &key, std::string uniqueName = buildUniqueFunctionName(key, selectedName.getValue()); - if (auto existing = mod.lookupSymbol(uniqueName)) + if (auto existing = mod.lookupSymbol(uniqueName)) { return existing; + } return invokeInProcessTileLib(key, selectedName.getValue(), uniqueName, mod, ctx); @@ -1169,8 +1279,9 @@ LogicalResult ExpandState::expandTileOpsInFunction(func::FuncOp func, // Collect tile ops first (avoid modifying while iterating). SmallVector tileOps; func.walk([&](Operation *op) { - if (pto::isTileLibExpandableOp(op)) + if (pto::isTileLibExpandableOp(op)) { tileOps.push_back(op); + } }); for (auto *op : tileOps) { @@ -1224,8 +1335,9 @@ void ExpandTileOpPass::runOnOperation() { } return WalkResult::advance(); }); - if (!hasExpandableOps) + if (!hasExpandableOps) { return; + } std::shared_ptr tileLibService = pto::TileLibRuntime::getService(); @@ -1239,10 +1351,12 @@ void ExpandTileOpPass::runOnOperation() { state.tileLibService = tileLibService; for (auto func : mod.getOps()) { - if (func.isExternal()) + if (func.isExternal()) { continue; - if (failed(state.expandTileOpsInFunction(func, mod, ctx))) + } + if (failed(state.expandTileOpsInFunction(func, mod, ctx))) { return signalPassFailure(); + } } } diff --git a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp index 593d234658..46902ed380 100644 --- a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp +++ b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp @@ -65,12 +65,15 @@ constexpr llvm::StringLiteral kTileOpValidShapeReadAttr = "__pto.tileop_valid_shape_abi"; static FailureOr parseFoldIntrinsicMode(StringRef mode) { - if (mode.empty() || mode == "all") + if (mode.empty() || mode == "all") { return FoldIntrinsicMode::All; - if (mode == "shape-only") + } + if (mode == "shape-only") { return FoldIntrinsicMode::ShapeOnly; - if (mode == "addr-only") + } + if (mode == "addr-only") { return FoldIntrinsicMode::AddrOnly; + } return failure(); } @@ -85,12 +88,14 @@ static bool shouldFoldAddrFamily(FoldIntrinsicMode mode) { static bool eraseDeadAllocTileOps(func::FuncOp func) { SmallVector deadAllocs; func.walk([&](pto::AllocTileOp alloc) { - if (alloc.getResult().use_empty()) + if (alloc.getResult().use_empty()) { deadAllocs.push_back(alloc); + } }); - for (pto::AllocTileOp alloc : llvm::reverse(deadAllocs)) + for (pto::AllocTileOp alloc : llvm::reverse(deadAllocs)) { alloc.erase(); + } return !deadAllocs.empty(); } @@ -113,8 +118,9 @@ static bool eraseDeadViewBridgeCasts(func::FuncOp func) { deadCasts.push_back(castOp); }); - for (auto castOp : llvm::reverse(deadCasts)) + for (auto castOp : llvm::reverse(deadCasts)) { castOp.erase(); + } return !deadCasts.empty(); } @@ -127,8 +133,9 @@ static bool eraseDeadMemrefViewOps(func::FuncOp func) { deadMemrefOps.push_back(op); }); - for (Operation *op : llvm::reverse(deadMemrefOps)) + for (Operation *op : llvm::reverse(deadMemrefOps)) { op->erase(); + } return !deadMemrefOps.empty(); } @@ -141,8 +148,9 @@ static bool eraseDeadTensorViewOps(func::FuncOp func) { deadViewOps.push_back(op); }); - for (Operation *op : llvm::reverse(deadViewOps)) + for (Operation *op : llvm::reverse(deadViewOps)) { op->erase(); + } return !deadViewOps.empty(); } @@ -162,8 +170,9 @@ struct TileHandleInfo { static std::pair findSetValidShapeOverride(Value tileBuf) { for (Operation *user : tileBuf.getUsers()) { auto setValid = dyn_cast(user); - if (!setValid || setValid.getSource() != tileBuf) + if (!setValid || setValid.getSource() != tileBuf) { continue; + } return {setValid.getValidRow(), setValid.getValidCol()}; } return {Value(), Value()}; @@ -176,8 +185,9 @@ static std::pair findSetValidShapeOverride(Value tileBuf) { static Value unwrapBridgingCasts(Value v) { while (v) { Operation *defOp = v.getDefiningOp(); - if (!defOp) + if (!defOp) { break; + } if (auto cast = dyn_cast(defOp)) { if (cast.getNumOperands() == 1 && cast.getNumResults() == 1) { v = cast.getOperand(0); @@ -244,8 +254,9 @@ static std::optional resolveTileHandle(Value tileBuf, if (auto reshape = tileBuf.getDefiningOp()) { auto sourceInfo = resolveTileHandle(reshape.getSrc(), user); - if (!sourceInfo) + if (!sourceInfo) { return std::nullopt; + } auto tileTy = dyn_cast(reshape.getResult().getType()); if (!tileTy) { @@ -348,21 +359,26 @@ static bool getConstIndexValue(Value v, int64_t &out) { return true; } } - if (auto castOp = v.getDefiningOp()) + if (auto castOp = v.getDefiningOp()) { return getConstIndexValue(castOp.getIn(), out); - if (auto extOp = v.getDefiningOp()) + } + if (auto extOp = v.getDefiningOp()) { return getConstIndexValue(extOp.getIn(), out); - if (auto extOp = v.getDefiningOp()) + } + if (auto extOp = v.getDefiningOp()) { return getConstIndexValue(extOp.getIn(), out); - if (auto truncOp = v.getDefiningOp()) + } + if (auto truncOp = v.getDefiningOp()) { return getConstIndexValue(truncOp.getIn(), out); + } return false; } static Value getValueOrCreateConstant(OpBuilder &builder, Location loc, OpFoldResult ofr) { - if (auto val = dyn_cast(ofr)) + if (auto val = dyn_cast(ofr)) { return val; + } auto intAttr = dyn_cast(cast(ofr)); assert(intAttr && "expected integer attribute in OpFoldResult"); return builder.create(loc, intAttr.getInt()); @@ -371,11 +387,13 @@ static Value getValueOrCreateConstant(OpBuilder &builder, Location loc, static bool isAllStaticZero(ArrayRef ofrs) { for (OpFoldResult ofr : ofrs) { auto attr = dyn_cast(ofr); - if (!attr) + if (!attr) { return false; + } auto intAttr = dyn_cast(attr); - if (!intAttr || intAttr.getInt() != 0) + if (!intAttr || intAttr.getInt() != 0) { return false; + } } return true; } @@ -385,8 +403,9 @@ static Value computeResultStride(OpBuilder &builder, Location loc, OpFoldResult svStride) { if (auto attr = dyn_cast(svStride)) { auto intAttr = dyn_cast(attr); - if (intAttr && intAttr.getInt() == 1) + if (intAttr && intAttr.getInt() == 1) { return getValueOrCreateConstant(builder, loc, rcStride); + } } Value lhs = getValueOrCreateConstant(builder, loc, rcStride); @@ -401,16 +420,18 @@ static Value computeLinearOffset(OpBuilder &builder, Location loc, bool rcAllZero = isAllStaticZero(rcOffsets); bool svAllZero = isAllStaticZero(svOffsets); - if (rcAllZero && svAllZero) + if (rcAllZero && svAllZero) { return Value(); + } Value svPart; if (!svAllZero) { for (auto [svOffset, rcStride] : llvm::zip(svOffsets, rcStrides)) { if (auto attr = dyn_cast(svOffset)) { auto intAttr = dyn_cast(attr); - if (intAttr && intAttr.getInt() == 0) + if (intAttr && intAttr.getInt() == 0) { continue; + } } Value off = getValueOrCreateConstant(builder, loc, svOffset); @@ -422,20 +443,23 @@ static Value computeLinearOffset(OpBuilder &builder, Location loc, Value rcPart; if (!rcAllZero) { - if (rcOffsets.empty()) + if (rcOffsets.empty()) { return Value(); + } rcPart = getValueOrCreateConstant(builder, loc, rcOffsets.front()); } - if (rcPart && svPart) + if (rcPart && svPart) { return builder.create(loc, rcPart, svPart); + } return rcPart ? rcPart : svPart; } static Value unwrapPTOViewBridge(Value value) { while (auto cast = value.getDefiningOp()) { - if (cast.getNumOperands() != 1 || cast.getNumResults() != 1) + if (cast.getNumOperands() != 1 || cast.getNumResults() != 1) { break; + } value = cast.getOperand(0); } return value; @@ -474,8 +498,9 @@ static Value resolvePTOViewStride(Value view, int64_t dim, OpBuilder &builder, if (Value projected = projectSCFIfViewResult( view, PTOViewProjectionKind::Stride, dim, {}, builder, user)) return projected; - if (auto partition = view.getDefiningOp()) + if (auto partition = view.getDefiningOp()) { return resolvePTOViewStride(partition.getSource(), dim, builder, user); + } if (auto makeView = view.getDefiningOp()) { if (dim < 0 || dim >= static_cast(makeView.getStrides().size())) return {}; @@ -492,8 +517,9 @@ static Value resolvePTOViewAddress(Value view, pto::PtrType resultType, return projected; if (auto makeView = view.getDefiningOp()) { Value ptr = makeView.getPtr(); - if (ptr.getType() == resultType) + if (ptr.getType() == resultType) { return ptr; + } return builder.create(user->getLoc(), resultType, ptr); } auto partition = view.getDefiningOp(); @@ -517,8 +543,9 @@ static Value resolvePTOViewAddress(Value view, pto::PtrType resultType, resolvePTOViewAddress(partition.getSource(), resultType, builder, user); if (!base) return {}; - if (!linearOffset) + if (!linearOffset) { return base; + } return builder.create(user->getLoc(), resultType, base, linearOffset); } @@ -527,8 +554,9 @@ static void cloneBlockWithoutTerminator(Block *source, Block *target, IRMapping &mapping, OpBuilder &builder) { builder.setInsertionPointToStart(target); - for (Operation &op : source->without_terminator()) + for (Operation &op : source->without_terminator()) { builder.clone(op, mapping); + } } static Value projectSCFIfViewResult(Value view, PTOViewProjectionKind kind, @@ -569,8 +597,9 @@ static Value projectSCFIfViewResult(Value view, PTOViewProjectionKind kind, cloneBlockWithoutTerminator(source, target, mapping, ifBuilder); SmallVector yields; yields.reserve(oldYield.getNumOperands() + 1); - for (Value operand : oldYield.getOperands()) + for (Value operand : oldYield.getOperands()) { yields.push_back(mapping.lookupOrDefault(operand)); + } Value yieldedView = yields[resultIndex]; ifBuilder.setInsertionPointToEnd(target); Value projection; @@ -635,8 +664,9 @@ struct FoldTileBufIntrinsicsPass // ops on tile_buf function arguments — they have no materialized tile // handle anchor to fold against and will be removed by later DCE. Skip // them. - if (func->hasAttr("pto.tilelang.instance")) + if (func->hasAttr("pto.tilelang.instance")) { return; + } SmallVector addrOps; SmallVector rowsOps; @@ -647,8 +677,9 @@ struct FoldTileBufIntrinsicsPass SmallVector getValidShapeOps; func.walk([&](Operation *op) { - if (auto addr = dyn_cast(op)) + if (auto addr = dyn_cast(op)) { addrOps.push_back(addr); + } else if (auto rows = dyn_cast(op)) rowsOps.push_back(rows); else if (auto cols = dyn_cast(op)) @@ -671,10 +702,12 @@ struct FoldTileBufIntrinsicsPass // resolveTileHandle observe the overridden valid shape carried by a // treshape + set_validshape pair. for (auto gvsOp : getValidShapeOps) { - if (gvsOp->hasAttr(kTileOpValidShapeReadAttr)) + if (gvsOp->hasAttr(kTileOpValidShapeReadAttr)) { continue; - if (!isa(gvsOp.getSource().getType())) + } + if (!isa(gvsOp.getSource().getType())) { continue; + } builder.setInsertionPoint(gvsOp); auto tileTy = cast(gvsOp.getSource().getType()); @@ -692,12 +725,15 @@ struct FoldTileBufIntrinsicsPass if (!rowReplacement || !colReplacement) { auto handleInfo = resolveTileHandle(gvsOp.getSource(), gvsOp); - if (!handleInfo) + if (!handleInfo) { return signalPassFailure(); - if (!rowReplacement) + } + if (!rowReplacement) { rowReplacement = handleInfo->validRow; - if (!colReplacement) + } + if (!colReplacement) { colReplacement = handleInfo->validCol; + } } if (!rowReplacement || !colReplacement) { @@ -723,8 +759,9 @@ struct FoldTileBufIntrinsicsPass dyn_cast(addrOp.getSrc().getType())) { if (auto resultMemrefType = dyn_cast(addrOp.getDst().getType())) { - if (srcMemrefType != resultMemrefType) + if (srcMemrefType != resultMemrefType) { addrOp.getDst().setType(srcMemrefType); + } addrOp.getDst().replaceAllUsesWith(addrOp.getSrc()); addrOp.erase(); continue; @@ -749,12 +786,14 @@ struct FoldTileBufIntrinsicsPass // Keep tile_buf_addr attached to that handle; VPTO pointer // normalization converts it directly without choosing one branch's // allocation address here. - if (isSCFTileCarrier(addrOp.getSrc())) + if (isSCFTileCarrier(addrOp.getSrc())) { continue; + } auto handleInfo = resolveTileHandle(addrOp.getSrc(), addrOp); - if (!handleInfo) + if (!handleInfo) { return signalPassFailure(); + } auto tileTy = dyn_cast(addrOp.getSrc().getType()); if (!tileTy) { @@ -803,8 +842,9 @@ struct FoldTileBufIntrinsicsPass builder.create(rowsOp.getLoc(), vRow); } else { auto handleInfo = resolveTileHandle(rowsOp.getSrc(), rowsOp); - if (!handleInfo) + if (!handleInfo) { return signalPassFailure(); + } replacement = handleInfo->validRow; if (!replacement) { rowsOp.emitError( @@ -836,8 +876,9 @@ struct FoldTileBufIntrinsicsPass builder.create(colsOp.getLoc(), vCol); } else { auto handleInfo = resolveTileHandle(colsOp.getSrc(), colsOp); - if (!handleInfo) + if (!handleInfo) { return signalPassFailure(); + } replacement = handleInfo->validCol; if (!replacement) { colsOp.emitError( @@ -870,8 +911,9 @@ struct FoldTileBufIntrinsicsPass } auto chain = traceViewChain(dimOp.getTensorView(), dimOp); - if (!chain) + if (!chain) { return signalPassFailure(); + } auto svTy = cast(chain->subview.getType()); if (dimIdx < 0 || dimIdx >= svTy.getRank()) { @@ -914,8 +956,9 @@ struct FoldTileBufIntrinsicsPass } auto chain = traceViewChain(strideOp.getTensorView(), strideOp); - if (!chain) + if (!chain) { return signalPassFailure(); + } auto svTy = cast(chain->subview.getType()); if (dimIdx < 0 || dimIdx >= svTy.getRank()) { @@ -951,15 +994,17 @@ struct FoldTileBufIntrinsicsPass } auto chain = traceViewChain(addrOp.getSrc(), addrOp); - if (!chain) + if (!chain) { return signalPassFailure(); + } if (!resultPtrType) { if (auto resultMemrefType = dyn_cast(addrOp.getDst().getType())) { Value base = chain->baseMemref; - if (base.getType() != resultMemrefType) + if (base.getType() != resultMemrefType) { addrOp.getDst().setType(cast(base.getType())); + } addrOp.getDst().replaceAllUsesWith(base); addrOp.erase(); continue; @@ -1000,8 +1045,9 @@ struct FoldTileBufIntrinsicsPass castOp.getResult(0).getType())) deadCasts.push_back(castOp); }); - for (auto castOp : llvm::reverse(deadCasts)) + for (auto castOp : llvm::reverse(deadCasts)) { castOp.erase(); + } while (true) { SmallVector deadMemrefOps; @@ -1011,10 +1057,12 @@ struct FoldTileBufIntrinsicsPass op->use_empty()) deadMemrefOps.push_back(op); }); - if (deadMemrefOps.empty()) + if (deadMemrefOps.empty()) { break; - for (auto *op : llvm::reverse(deadMemrefOps)) + } + for (auto *op : llvm::reverse(deadMemrefOps)) { op->erase(); + } } // Erase metadata writes only after every reader has been folded. TileOp @@ -1032,8 +1080,9 @@ struct FoldTileBufIntrinsicsPass } return WalkResult::advance(); }); - if (!hasRuntimeReader) + if (!hasRuntimeReader) { op.erase(); + } } // DCE tile-handle view / alloc ops left behind after valid-shape @@ -1043,10 +1092,12 @@ struct FoldTileBufIntrinsicsPass tileDceChanged = false; SmallVector deadTileOps; func.walk([&](Operation *op) { - if (!op->use_empty()) + if (!op->use_empty()) { return; - if (isa(op)) + } + if (isa(op)) { deadTileOps.push_back(op); + } else if (auto castOp = dyn_cast(op)) { if (castOp.getNumOperands() == 1 && isa(castOp.getResult(0).getType())) diff --git a/lib/PTO/Transforms/GraphSyncSolver/Utility.cpp b/lib/PTO/Transforms/GraphSyncSolver/Utility.cpp index 20fbd5c7ff..fe5185e2b1 100644 --- a/lib/PTO/Transforms/GraphSyncSolver/Utility.cpp +++ b/lib/PTO/Transforms/GraphSyncSolver/Utility.cpp @@ -257,7 +257,6 @@ namespace mlir::pto::syncsolver { // Check if two integer ranges intersect (half-open semantics: [l, r) ) bool checkRangesIntersect(int l1, int r1, int l2, int r2) { - // return !(r1 <= l2 || r2 <= l1); return r1 > l2 && r2 > l1; } diff --git a/lib/PTO/Transforms/InferPTOLayout.cpp b/lib/PTO/Transforms/InferPTOLayout.cpp index 572995a870..77194afd6a 100644 --- a/lib/PTO/Transforms/InferPTOLayout.cpp +++ b/lib/PTO/Transforms/InferPTOLayout.cpp @@ -52,13 +52,16 @@ static constexpr int64_t kNZFractalBytes = 512; using LayoutRankVector = SmallVector; static std::optional getConstInt(Value v) { - if (auto c = v.getDefiningOp()) + if (auto c = v.getDefiningOp()) { return c.value(); - if (auto c = v.getDefiningOp()) + } + if (auto c = v.getDefiningOp()) { return c.value(); + } if (auto c = v.getDefiningOp()) { - if (auto ia = dyn_cast(c.getValue())) + if (auto ia = dyn_cast(c.getValue())) { return ia.getInt(); + } } return std::nullopt; } @@ -66,8 +69,9 @@ static std::optional getConstInt(Value v) { static std::optional getConstInt(OpFoldResult ofr) { if (isa(ofr)) { Attribute attr = cast(ofr); - if (auto ia = dyn_cast(attr)) + if (auto ia = dyn_cast(attr)) { return ia.getInt(); + } return std::nullopt; } return getConstInt(cast(ofr)); @@ -110,10 +114,12 @@ static bool isMinor2DLayout(Layout layout) { static std::optional rightAlignTo5D(ArrayRef shape, ArrayRef stride) { - if (shape.size() != stride.size()) + if (shape.size() != stride.size()) { return std::nullopt; - if (shape.size() > kPaddedLayoutRank) + } + if (shape.size() > kPaddedLayoutRank) { return std::nullopt; + } ShapeStride5D out; out.shape.assign(kPaddedLayoutRank, kUnitExtent); @@ -128,27 +134,32 @@ static std::optional rightAlignTo5D(ArrayRef shape, // Derive the padded leading strides with the same rule used in EmitC: // stride[i] = shape[i+1] * stride[i+1]. - for (int i = shift - 1; i >= 0; --i) + for (int i = shift - 1; i >= 0; --i) { out.stride[i] = out.shape[i + 1] * out.stride[i + 1]; + } return out; } static bool matchesNDMinor2D(int64_t rows, int64_t cols, int64_t rowStride, int64_t colStride) { - if (cols != 1 && colStride != 1) + if (cols != 1 && colStride != 1) { return false; - if (rows == 1) + } + if (rows == 1) { return true; + } return cols == 1 ? rowStride == 1 : rowStride == cols; } static bool matchesDNMinor2D(int64_t rows, int64_t cols, int64_t rowStride, int64_t colStride) { - if (rows != 1 && rowStride != 1) + if (rows != 1 && rowStride != 1) { return false; - if (cols == 1) + } + if (cols == 1) { return true; + } return rows == 1 ? colStride == 1 : colStride == rows; } @@ -157,11 +168,13 @@ static std::optional inferMinor2DLayout( std::optional preferredMinor2D, bool *isMinor2DAmbiguous) { const bool nd = matchesNDMinor2D(rows, cols, rowStride, colStride); const bool dn = matchesDNMinor2D(rows, cols, rowStride, colStride); - if (!nd && !dn) + if (!nd && !dn) { return Layout::ND; + } if (nd && dn) { - if (isMinor2DAmbiguous) + if (isMinor2DAmbiguous) { *isMinor2DAmbiguous = true; + } if (preferredMinor2D && (*preferredMinor2D == Layout::ND || *preferredMinor2D == Layout::DN)) { return *preferredMinor2D; @@ -183,8 +196,9 @@ static std::optional inferNZLayout(ArrayRef shape, (sh3 == kNZInnerRows) && (sh3 * sh4 * static_cast(elemBytes) == kNZFractalBytes); bool strideMatch = (st5 == kUnitExtent) && (st4 == sh5); - if (alignMatch && strideMatch) + if (alignMatch && strideMatch) { return Layout::NZ; + } return std::nullopt; } @@ -194,16 +208,20 @@ static std::optional inferLayout5D(ArrayRef shape, std::optional preferredMinor2D = std::nullopt, bool *isMinor2DAmbiguous = nullptr) { - if (shape.size() != strides.size() || elemBytes == 0) + if (shape.size() != strides.size() || elemBytes == 0) { return std::nullopt; - if (isMinor2DAmbiguous) + } + if (isMinor2DAmbiguous) { *isMinor2DAmbiguous = false; + } auto padded = rightAlignTo5D(shape, strides); - if (!padded) + if (!padded) { return std::nullopt; + } - if (auto nz = inferNZLayout(padded->shape, padded->stride, elemBytes)) + if (auto nz = inferNZLayout(padded->shape, padded->stride, elemBytes)) { return nz; + } const int64_t rows = padded->shape[3]; const int64_t cols = padded->shape[4]; @@ -215,11 +233,13 @@ static std::optional inferLayout5D(ArrayRef shape, static std::optional tileBLayoutToGlobalLayout(Type tileLikeTy) { auto tbTy = dyn_cast(tileLikeTy); - if (!tbTy) + if (!tbTy) { return std::nullopt; + } auto bl = dyn_cast_or_null(tbTy.getBLayoutAttr()); - if (!bl) + if (!bl) { return std::nullopt; + } switch (bl.getValue()) { case BLayout::RowMajor: return Layout::ND; @@ -231,8 +251,9 @@ static std::optional tileBLayoutToGlobalLayout(Type tileLikeTy) { static bool isVectorTileType(Type tileLikeTy) { auto tbTy = dyn_cast(tileLikeTy); - if (!tbTy) + if (!tbTy) { return false; + } auto ms = dyn_cast_or_null(tbTy.getMemorySpace()); return ms && ms.getAddressSpace() == AddressSpace::VEC; } @@ -254,10 +275,11 @@ static ResolvedLayoutInfo resolveLayoutFromViewValue(Value v); static void setLayoutAttr(Operation *op, Layout layout, bool inferred) { op->setAttr(kLayoutAttrName, LayoutAttr::get(op->getContext(), layout)); - if (inferred) + if (inferred) { op->setAttr(kInferredLayoutAttrName, BoolAttr::get(op->getContext(), true)); - else + } else { op->removeAttr(kInferredLayoutAttrName); + } } template @@ -306,8 +328,9 @@ static void maybeRepairMinor2DLoadStoreLayout(LoadStoreOp op, ViewGetter getView auto tilePref = isVectorTileType(getTile(op).getType()) ? tileBLayoutToGlobalLayout(getTile(op).getType()) : std::nullopt; - if (!tilePref || (*tilePref != Layout::ND && *tilePref != Layout::DN)) + if (!tilePref || (*tilePref != Layout::ND && *tilePref != Layout::DN)) { return; + } auto viewInfo = resolveLayoutFromViewValue(getView(op)); if (!viewInfo.owner || !viewInfo.layout || !viewInfo.inferred || @@ -315,13 +338,15 @@ static void maybeRepairMinor2DLoadStoreLayout(LoadStoreOp op, ViewGetter getView return; } auto tv = dyn_cast(viewInfo.owner); - if (!tv) + if (!tv) { return; + } SmallVector shape, strides; bool ambiguous = false; - if (!getStaticShapeAndStride(tv, shape, strides)) + if (!getStaticShapeAndStride(tv, shape, strides)) { return; + } (void)inferLayout5D( shape, strides, elemByteSize(cast(tv.getResult().getType()).getElementType()), @@ -360,8 +385,9 @@ struct LayoutPreference { static LayoutPreference collectPreferredLayoutFromConsumers(Value tensorView) { LayoutPreference result; auto mergePref = [&](std::optional candidate) { - if (!candidate) + if (!candidate) { return; + } if (!result.preferred) { result.preferred = candidate; return; @@ -378,8 +404,9 @@ static LayoutPreference collectPreferredLayoutFromConsumers(Value tensorView) { unsigned operandIndex = use.getOperandNumber(); if (auto part = dyn_cast(owner)) { - if (operandIndex == 0) + if (operandIndex == 0) { self(self, part.getResult()); + } continue; } @@ -413,8 +440,9 @@ static LayoutPreference collectPreferredLayoutFromConsumers(Value tensorView) { } if (auto store = dyn_cast(owner)) { - if (operandIndex == 1 && isVectorTileType(store.getSrc().getType())) + if (operandIndex == 1 && isVectorTileType(store.getSrc().getType())) { mergePref(tileBLayoutToGlobalLayout(store.getSrc().getType())); + } continue; } } @@ -433,8 +461,9 @@ static std::optional inferMakeTensorViewLayout( *pref.preferred == Layout::MX_B_NN)) return pref.preferred; std::optional preferredForAmbiguous = std::nullopt; - if (!pref.conflict && isMinorColsOne(shape)) + if (!pref.conflict && isMinorColsOne(shape)) { preferredForAmbiguous = pref.preferred; + } return inferLayout5D( shape, strides, elemByteSize( @@ -445,25 +474,30 @@ static std::optional inferMakeTensorViewLayout( static void reconcileAmbiguousTensorViewLayout(MakeTensorViewOp op, ArrayRef shape) { auto pref = collectPreferredLayoutFromConsumers(op.getResult()); - if (!isMinorColsOne(shape)) + if (!isMinorColsOne(shape)) { return; - if (!op->getAttrOfType(kInferredLayoutAttrName)) + } + if (!op->getAttrOfType(kInferredLayoutAttrName)) { return; + } auto cur = op->getAttrOfType(kLayoutAttrName); - if (cur && pref.preferred && *pref.preferred != cur.getLayout()) + if (cur && pref.preferred && *pref.preferred != cur.getLayout()) { setLayoutAttr(op.getOperation(), *pref.preferred, /*inferred=*/true); + } } static bool getStaticShapeAndStride(MakeTensorViewOp op, SmallVectorImpl &shape, SmallVectorImpl &strides) { auto tvTy = dyn_cast(op.getResult().getType()); - if (!tvTy) + if (!tvTy) { return false; + } const size_t rank = op.getShape().size(); - if (rank == 0 || rank > kPaddedLayoutRank) + if (rank == 0 || rank > kPaddedLayoutRank) { return false; + } shape.clear(); shape.reserve(rank); @@ -471,8 +505,9 @@ static bool getStaticShapeAndStride(MakeTensorViewOp op, int64_t dim = tvTy.getShape()[i]; if (dim == ShapedType::kDynamic) { auto v = getConstInt(op.getShape()[i]); - if (!v) + if (!v) { return false; + } dim = *v; } shape.push_back(dim); @@ -482,8 +517,9 @@ static bool getStaticShapeAndStride(MakeTensorViewOp op, strides.reserve(rank); for (Value s : op.getStrides()) { auto v = getConstInt(s); - if (!v) + if (!v) { return false; + } strides.push_back(*v); } return true; @@ -543,16 +579,18 @@ static void inferMakeTensorViewLayoutAttr(MakeTensorViewOp op, auto inferred = inferMakeTensorViewLayout(op, shape, strides, isAmbiguous); verifyOrSetLayoutAttr(op.getOperation(), inferred, signalFailure, isAmbiguous); - if (isAmbiguous) + if (isAmbiguous) { reconcileAmbiguousTensorViewLayout(op, shape); + } } template static void inferReinterpretCastLayoutAttr(memref::ReinterpretCastOp op, SignalFailureFn signalFailure) { auto mrTy = dyn_cast(op.getType()); - if (!mrTy || !isGlobalMemRef(mrTy)) + if (!mrTy || !isGlobalMemRef(mrTy)) { return; + } const size_t rank = op.getMixedSizes().size(); if (rank == 0 || rank > kPaddedLayoutRank) { @@ -608,11 +646,13 @@ struct InferPTOLayoutPass // ------------------------------------------------------------------ func.walk([&](memref::SubViewOp op) { auto resTy = dyn_cast(op.getType()); - if (!resTy || !isGlobalMemRef(resTy)) + if (!resTy || !isGlobalMemRef(resTy)) { return; + } - if (op->getAttrOfType(kLayoutAttrName)) + if (op->getAttrOfType(kLayoutAttrName)) { return; + } if (Operation *def = op.getSource().getDefiningOp()) { if (auto srcLayout = def->getAttrOfType(kLayoutAttrName)) { diff --git a/lib/PTO/Transforms/InferPTOMemScope.cpp b/lib/PTO/Transforms/InferPTOMemScope.cpp index f62eee2174..79681e1130 100644 --- a/lib/PTO/Transforms/InferPTOMemScope.cpp +++ b/lib/PTO/Transforms/InferPTOMemScope.cpp @@ -39,9 +39,10 @@ namespace { static std::optional requireRootAlloc(Operation *op, Value value, StringRef valueName) { auto alloc = tracebackMemRefToAlloc(value); - if (!alloc.has_value()) + if (!alloc.has_value()) { emitError(op->getLoc()) << "Cannot find root memref.alloc for " << valueName << " of this op."; + } return alloc; } @@ -50,11 +51,13 @@ static LogicalResult propagateAllocScope(Operation *op, Value value, const AddressSpaceAttr &targetScope, MemScopeInferAndPropagateHelper &helper) { auto alloc = requireRootAlloc(op, value, valueName); - if (!alloc.has_value()) + if (!alloc.has_value()) { return failure(); - if (failed(helper.Run(*alloc, targetScope))) + } + if (failed(helper.Run(*alloc, targetScope))) { return op->emitOpError() << "Failed to infer/propagate memory scope for " << valueName; + } return success(); } @@ -117,8 +120,9 @@ propagateMemScopeToUser(MemScopeInferAndPropagateHelper &helper, Value val, }) .Case([&](auto) { return success(); }) .Default([&](Operation *op) { - if (op->getNumResults() == 0 || !hasMemRefResults(op)) + if (op->getNumResults() == 0 || !hasMemRefResults(op)) { return success(); + } op->emitOpError("Unsupported user for root alloc op."); return failure(); }); @@ -189,27 +193,32 @@ static LogicalResult propagateOperandScopes( ArrayRef> specs) { MemScopeInferAndPropagateHelper helper; for (const auto &[value, valueName, targetScope] : specs) { - if (failed(propagateAllocScope(op, value, valueName, targetScope, helper))) + if (failed(propagateAllocScope(op, value, valueName, targetScope, helper))) { return failure(); + } } return success(); } LogicalResult pto::inferAndPropagateMemScopeForMovDps(pto::TMovOp op) { - if (failed(ensureDpsOnlyOp(op))) + if (failed(ensureDpsOnlyOp(op))) { return failure(); + } auto dstAlloc = requireRootAlloc(op, op.getDst(), "mB"); - if (!dstAlloc.has_value()) + if (!dstAlloc.has_value()) { return failure(); + } auto memRefType = dyn_cast(dstAlloc->getType()); - if (!memRefType) + if (!memRefType) { return op->emitOpError("Failed to infer/propagate memory scope for mA"); + } auto memSpace = memRefType.getMemorySpace(); - if (!memSpace) + if (!memSpace) { return success(); + } auto l0aSpaceAttr = getMemScopeAttr(op->getContext(), pto::AddressSpace::LEFT); @@ -221,10 +230,12 @@ LogicalResult pto::inferAndPropagateMemScopeForMovDps(pto::TMovOp op) { auto ubSpaceAttr = getMemScopeAttr(op->getContext(), pto::AddressSpace::VEC); auto biasSpaceAttr = getMemScopeAttr(op->getContext(), pto::AddressSpace::BIAS); - if (memSpace == ubSpaceAttr) + if (memSpace == ubSpaceAttr) { return propagateOperandScopes(op, {{op.getSrc(), "mA", ubSpaceAttr}}); - if (memSpace == l1SpaceAttr) + } + if (memSpace == l1SpaceAttr) { return propagateOperandScopes(op, {{op.getSrc(), "mA", l0cSpaceAttr}}); + } if (memSpace == l0aSpaceAttr || memSpace == l0bSpaceAttr || memSpace == biasSpaceAttr) { return propagateOperandScopes(op, {{op.getSrc(), "mA", l1SpaceAttr}}); @@ -233,8 +244,9 @@ LogicalResult pto::inferAndPropagateMemScopeForMovDps(pto::TMovOp op) { } LogicalResult pto::inferAndPropagateMemScopeForMatmulAccDps(pto::TMatmulAccOp op) { - if (failed(ensureDpsOnlyOp(op))) + if (failed(ensureDpsOnlyOp(op))) { return failure(); + } return propagateOperandScopes( op, {{op.getAccIn(), "mAcc", @@ -249,8 +261,9 @@ LogicalResult pto::inferAndPropagateMemScopeForMatmulAccDps(pto::TMatmulAccOp op LogicalResult pto::inferAndPropagateMemScopeForMatmulBiasDps(pto::TMatmulBiasOp op) { - if (failed(ensureDpsOnlyOp(op))) + if (failed(ensureDpsOnlyOp(op))) { return failure(); + } return propagateOperandScopes( op, {{op.getA(), "mA", @@ -264,8 +277,9 @@ LogicalResult pto::inferAndPropagateMemScopeForMatmulBiasDps(pto::TMatmulBiasOp } LogicalResult pto::inferAndPropagateMemScopeForMatmulDps(pto::TMatmulOp op) { - if (failed(ensureDpsOnlyOp(op))) + if (failed(ensureDpsOnlyOp(op))) { return failure(); + } return propagateOperandScopes( op, {{op.getLhs(), "mA", @@ -284,12 +298,14 @@ LogicalResult InferPTOMemScopePass::fixDeviceCallSite(func::FuncOp op) { func::CallOp call = cast(use.getUser()); // propagate call operand's memory scope for (auto [idx, callOperand] : llvm::enumerate(call.getArgOperands())) { - if (!isa(callOperand.getType())) + if (!isa(callOperand.getType())) { continue; + } auto funcOperandType = op.getFunctionType().getInput(idx); - if (!isa(funcOperandType)) + if (!isa(funcOperandType)) { continue; + } LDBG("call operand: " << callOperand); if (failed(helper.Run(tracebackMemRef(callOperand), @@ -302,12 +318,14 @@ LogicalResult InferPTOMemScopePass::fixDeviceCallSite(func::FuncOp op) { } // propagate call return value memory scope for (auto [idx, returnValue] : llvm::enumerate(call->getResults())) { - if (!isa(returnValue.getType())) + if (!isa(returnValue.getType())) { continue; + } auto funcReturnType = op.getFunctionType().getResult(idx); - if (!isa(funcReturnType)) + if (!isa(funcReturnType)) { continue; + } if (failed(helper.Run(returnValue, getPTOAddressSpaceAttr(funcReturnType)))) { @@ -326,12 +344,14 @@ LogicalResult InferPTOMemScopePass::fixDeviceCallSite(func::FuncOp op) { /// information to update the function's type. [[maybe_unused]] LogicalResult InferPTOMemScopePass::fixHostFuncSignature(func::FuncOp op) { // Skip external host functions because we know nothing about it. - if (op.isExternal()) + if (op.isExternal()) { return success(); + } func::ReturnOp returnOp = getAssumedUniqueReturnOp(op); - if (!returnOp) + if (!returnOp) { return failure(); + } SmallVector newArgsType(llvm::map_to_vector( op.getArguments(), [](const BlockArgument &ba) { return ba.getType(); })); @@ -342,9 +362,10 @@ LogicalResult InferPTOMemScopePass::fixDeviceCallSite(func::FuncOp op) { return success(); } -LogicalResult inferAndPropagateMemScopeForExternFunc(func::FuncOp op) { - if (!op.isExternal()) +static LogicalResult inferAndPropagateMemScopeForExternFunc(func::FuncOp op) { + if (!op.isExternal()) { return failure(); + } auto gmSpaceAttr = AddressSpaceAttr::get(op->getContext(), pto::AddressSpace::GM); @@ -354,8 +375,9 @@ LogicalResult inferAndPropagateMemScopeForExternFunc(func::FuncOp op) { for (auto &argType : newArgTypes) { // If not base memref and already has memspace then skip if (auto memrefType = dyn_cast(argType)) { - if (memrefType.getMemorySpace()) + if (memrefType.getMemorySpace()) { continue; + } argType = getBaseMemRefTypeWithNewScope(memrefType, gmSpaceAttr); } } @@ -365,8 +387,9 @@ LogicalResult inferAndPropagateMemScopeForExternFunc(func::FuncOp op) { for (auto &resultType : newReturnTypes) { // If not base memref and already has memspace then skip if (auto memrefType = dyn_cast(resultType)) { - if (memrefType.getMemorySpace()) + if (memrefType.getMemorySpace()) { continue; + } resultType = getBaseMemRefTypeWithNewScope(memrefType, gmSpaceAttr); } } @@ -376,8 +399,9 @@ LogicalResult inferAndPropagateMemScopeForExternFunc(func::FuncOp op) { } LogicalResult pto::inferAndPropagateMemScopeForFunc(func::FuncOp op) { - if (op.isExternal()) + if (op.isExternal()) { return inferAndPropagateMemScopeForExternFunc(op); + } LDBG("Begin infer and propagate memory scope for func" << op.getSymName()); MemScopeInferAndPropagateHelper helper; @@ -392,10 +416,11 @@ LogicalResult pto::inferAndPropagateMemScopeForFunc(func::FuncOp op) { } if (op->hasAttr(pto::VectorFunctionAttr::name)) { - if (failed(helper.Run(arg, ubSpaceAttr))) + if (failed(helper.Run(arg, ubSpaceAttr))) { return op->emitOpError() << "Failed to propagate UB memory scope for argument # in VF" << arg.getArgNumber(); + } } else if (failed(helper.Run(arg, gmSpaceAttr))) { return op->emitOpError() << "Failed to propagate memory scope for argument #" @@ -407,9 +432,10 @@ LogicalResult pto::inferAndPropagateMemScopeForFunc(func::FuncOp op) { op.getBody().front().getArgumentTypes(), op.getResultTypes()); op.setFunctionType(newFt); } - if (op->getNumResults() > 0) + if (op->getNumResults() > 0) { op.emitWarning() << "non-externl function has return value after bufferization!"; + } return success(); } @@ -446,8 +472,9 @@ LogicalResult pto::inferAndPropagateMemScopeForGpuFunc(gpu::GPUFuncOp op) { LogicalResult pto::inferAndPropagateUbufMemScope(memref::AllocOp op) { LDBG("Begin infer and propagate memory scope for: " << *op); auto memorySpace = op.getType().getMemorySpace(); - if (memorySpace) + if (memorySpace) { return success(); + } MemScopeInferAndPropagateHelper helper; auto ubSpaceAttr = @@ -473,47 +500,55 @@ void InferPTOMemScopePass::runOnOperation() { }); for (auto func : gpuFuncList) { - if (failed(inferAndPropagateMemScopeForGpuFunc(func))) + if (failed(inferAndPropagateMemScopeForGpuFunc(func))) { signalPassFailure(); + } } // Infer and propagate memory scope for device functions. for (auto func : deviceFuncList) { // Set the memory scope of values related to `pto::MmadL1Op` to L1 or L0C. func->walk([&](mlir::pto::TMatmulOp op) { - if (failed(pto::inferAndPropagateMemScopeForMatmulDps(op))) + if (failed(pto::inferAndPropagateMemScopeForMatmulDps(op))) { signalPassFailure(); + } }); func->walk([&](mlir::pto::TMatmulAccOp op) { - if (failed(pto::inferAndPropagateMemScopeForMatmulAccDps(op))) + if (failed(pto::inferAndPropagateMemScopeForMatmulAccDps(op))) { signalPassFailure(); + } }); func->walk([&](mlir::pto::TMatmulBiasOp op) { - if (failed(pto::inferAndPropagateMemScopeForMatmulBiasDps(op))) + if (failed(pto::inferAndPropagateMemScopeForMatmulBiasDps(op))) { signalPassFailure(); + } }); func->walk([&](mlir::pto::TMovOp op) { - if (failed(pto::inferAndPropagateMemScopeForMovDps(op))) + if (failed(pto::inferAndPropagateMemScopeForMovDps(op))) { signalPassFailure(); + } }); // Set device function arguments' memory scope to GM. - if (failed(pto::inferAndPropagateMemScopeForFunc(func))) + if (failed(pto::inferAndPropagateMemScopeForFunc(func))) { signalPassFailure(); + } // Finally, set the remaining memory scope in the device kernel to UB. func->walk([&](memref::AllocOp op) { - if (failed(pto::inferAndPropagateUbufMemScope(op))) + if (failed(pto::inferAndPropagateUbufMemScope(op))) { signalPassFailure(); + } }); } for (auto func : deviceFuncList) { - if (failed(fixDeviceCallSite(func))) + if (failed(fixDeviceCallSite(func))) { signalPassFailure(); + } } } diff --git a/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp b/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp index d2788df217..291f9500bd 100644 --- a/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp +++ b/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp @@ -180,8 +180,8 @@ bool MemoryDependentAnalyzer::MemAlias(const BaseMemInfo *a, llvm::errs() << " [MemAlias Check]\n"; printValueDebug(" Root A", a->rootBuffer); printValueDebug(" Root B", b->rootBuffer); - llvm::errs() << " Scope A: " << (int)as << ", Scope B: " << (int)bs - << "\n"; + llvm::errs() << " Scope A: " << static_cast(as) + << ", Scope B: " << static_cast(bs) << "\n"; } if (as != bs) { diff --git a/lib/PTO/Transforms/InsertSync/SyncCommon.cpp b/lib/PTO/Transforms/InsertSync/SyncCommon.cpp index e1efcf7d65..d58f39cf58 100644 --- a/lib/PTO/Transforms/InsertSync/SyncCommon.cpp +++ b/lib/PTO/Transforms/InsertSync/SyncCommon.cpp @@ -267,8 +267,8 @@ UNIT_FLAG CompoundInstanceElement::getUnitFlagMode() const { return it->second; } -Value getIsNotDeadLoopValue(scf::ForOp forOp, Location loc, - OpBuilder &rewriter) { +static Value getIsNotDeadLoopValue(scf::ForOp forOp, Location loc, + OpBuilder &rewriter) { Value upperBound = forOp.getUpperBound(); Value lowerBound = forOp.getLowerBound(); return rewriter.create(loc, arith::CmpIPredicate::slt, diff --git a/lib/PTO/Transforms/InsertTemplateAttributes.cpp b/lib/PTO/Transforms/InsertTemplateAttributes.cpp index eee8d3f879..0d1c68a4fc 100644 --- a/lib/PTO/Transforms/InsertTemplateAttributes.cpp +++ b/lib/PTO/Transforms/InsertTemplateAttributes.cpp @@ -55,50 +55,72 @@ struct CandidateMetadata { }; static std::string getDtypeString(Type elementType) { - if (elementType.isIndex()) + if (elementType.isIndex()) { return "i32"; - if (elementType.isInteger(1)) + } + if (elementType.isInteger(1)) { return "i1"; - if (elementType.isF32()) + } + if (elementType.isF32()) { return "f32"; - if (elementType.isF16()) + } + if (elementType.isF16()) { return "f16"; - if (elementType.isBF16()) + } + if (elementType.isBF16()) { return "bf16"; - if (isa(elementType)) + } + if (isa(elementType)) { return "f8e4m3"; - if (isa(elementType)) + } + if (isa(elementType)) { return "f8e5m2"; - if (isa(elementType)) + } + if (isa(elementType)) { return "hif8"; - if (isa(elementType)) + } + if (isa(elementType)) { return "f4e1m2x2"; - if (isa(elementType)) + } + if (isa(elementType)) { return "f4e2m1x2"; - if (elementType.isUnsignedInteger(64)) + } + if (elementType.isUnsignedInteger(64)) { return "ui64"; - if (elementType.isUnsignedInteger(32)) + } + if (elementType.isUnsignedInteger(32)) { return "ui32"; - if (elementType.isUnsignedInteger(16)) + } + if (elementType.isUnsignedInteger(16)) { return "ui16"; - if (elementType.isUnsignedInteger(8)) + } + if (elementType.isUnsignedInteger(8)) { return "ui8"; - if (elementType.isSignedInteger(64)) + } + if (elementType.isSignedInteger(64)) { return "si64"; - if (elementType.isSignedInteger(32)) + } + if (elementType.isSignedInteger(32)) { return "si32"; - if (elementType.isSignedInteger(16)) + } + if (elementType.isSignedInteger(16)) { return "si16"; - if (elementType.isSignedInteger(8)) + } + if (elementType.isSignedInteger(8)) { return "si8"; - if (elementType.isSignlessInteger(64)) + } + if (elementType.isSignlessInteger(64)) { return "i64"; - if (elementType.isSignlessInteger(32)) + } + if (elementType.isSignlessInteger(32)) { return "i32"; - if (elementType.isSignlessInteger(16)) + } + if (elementType.isSignlessInteger(16)) { return "i16"; - if (elementType.isSignlessInteger(8)) + } + if (elementType.isSignlessInteger(8)) { return "i8"; + } return ""; } @@ -152,18 +174,21 @@ static StringRef getBLayoutString(pto::BLayout layout) { } static StringRef getSLayoutString(pto::SLayout layout) { - if (layout == pto::SLayout::RowMajor) + if (layout == pto::SLayout::RowMajor) { return "row_major"; - if (layout == pto::SLayout::ColMajor) + } + if (layout == pto::SLayout::ColMajor) { return "col_major"; + } return "none_box"; } static void appendJsonIntArray(std::string &json, ArrayRef values) { json += "["; for (auto [index, value] : llvm::enumerate(values)) { - if (index != 0) + if (index != 0) { json += ","; + } json += std::to_string(value); } json += "]"; @@ -172,8 +197,9 @@ static void appendJsonIntArray(std::string &json, ArrayRef values) { static void appendJsonDimArray(std::string &json, ArrayRef values) { json += "["; for (auto [index, value] : llvm::enumerate(values)) { - if (index != 0) + if (index != 0) { json += ","; + } if (ShapedType::isDynamic(value)) { json += "null"; continue; @@ -198,14 +224,16 @@ static bool getStaticIntFromValue(Value value, int64_t &out) { static int64_t getStaticIntOrDynamic(OpFoldResult value) { if (isa(value)) { Attribute attr = cast(value); - if (auto integer = dyn_cast(attr)) + if (auto integer = dyn_cast(attr)) { return integer.getInt(); + } return ShapedType::kDynamic; } int64_t result = ShapedType::kDynamic; - if (getStaticIntFromValue(cast(value), result)) + if (getStaticIntFromValue(cast(value), result)) { return result; + } return ShapedType::kDynamic; } @@ -213,8 +241,9 @@ static void recordStaticSizes(ArrayRef values, SmallVectorImpl &out) { out.clear(); out.reserve(values.size()); - for (OpFoldResult value : values) + for (OpFoldResult value : values) { out.push_back(getStaticIntOrDynamic(value)); + } } static SmallVector @@ -237,16 +266,19 @@ combineSubviewStrides(ArrayRef baseStrides, static constexpr llvm::StringLiteral kLayoutAttrName = "layout"; static std::optional getLayoutAttrFromOp(Operation *op) { - if (!op) + if (!op) { return std::nullopt; - if (auto attr = op->getAttrOfType(kLayoutAttrName)) + } + if (auto attr = op->getAttrOfType(kLayoutAttrName)) { return attr.getLayout(); + } return std::nullopt; } static std::optional resolveViewLayout(Value value) { - if (!value) + if (!value) { return std::nullopt; + } Operation *definingOp = value.getDefiningOp(); while (definingOp) { @@ -255,8 +287,9 @@ static std::optional resolveViewLayout(Value value) { definingOp = value.getDefiningOp(); continue; } - if (auto layout = getLayoutAttrFromOp(definingOp)) + if (auto layout = getLayoutAttrFromOp(definingOp)) { return layout; + } if (auto subview = dyn_cast(definingOp)) { value = subview.getSource(); definingOp = value.getDefiningOp(); @@ -281,8 +314,9 @@ static std::optional resolveViewLayout(Value value) { static void populatePTOViewShapeAndStrides(Value value, SmallVectorImpl &shape, SmallVectorImpl &strides) { - if (!value) + if (!value) { return; + } if (auto part = value.getDefiningOp()) { if (shape.empty()) { @@ -295,24 +329,27 @@ static void populatePTOViewShapeAndStrides(Value value, if (shape.empty()) { auto partTy = dyn_cast(part.getResult().getType()); - if (partTy) + if (partTy) { shape.assign(partTy.getShape().begin(), partTy.getShape().end()); + } } } SmallVector sourceShape; SmallVector sourceStrides; populatePTOViewShapeAndStrides(part.getSource(), sourceShape, sourceStrides); - if (strides.empty() && !sourceStrides.empty()) + if (strides.empty() && !sourceStrides.empty()) { strides = sourceStrides; + } return; } if (auto make = value.getDefiningOp()) { if (shape.empty()) { auto viewTy = dyn_cast(make.getResult().getType()); - if (viewTy) + if (viewTy) { shape.assign(viewTy.getShape().begin(), viewTy.getShape().end()); + } } if (strides.empty()) { strides.reserve(make.getStrides().size()); @@ -326,25 +363,29 @@ static void populatePTOViewShapeAndStrides(Value value, } if (auto viewTy = dyn_cast(value.getType())) { - if (shape.empty()) + if (shape.empty()) { shape.assign(viewTy.getShape().begin(), viewTy.getShape().end()); + } } } static void populateViewShapeAndStrides(Value value, SmallVectorImpl &shape, SmallVectorImpl &strides) { - if (!value) + if (!value) { return; + } if (auto subview = value.getDefiningOp()) { populateViewShapeAndStrides(subview.getSource(), shape, strides); SmallVector subviewShape; recordStaticSizes(subview.getMixedSizes(), subviewShape); - if (!subviewShape.empty()) + if (!subviewShape.empty()) { shape = subviewShape; - if (!strides.empty()) + } + if (!strides.empty()) { strides = combineSubviewStrides(strides, subview.getMixedStrides()); + } return; } @@ -353,11 +394,13 @@ static void populateViewShapeAndStrides(Value value, if (shape.empty()) { SmallVector reinterpretShape; recordStaticSizes(reinterpret.getMixedSizes(), reinterpretShape); - if (!reinterpretShape.empty()) + if (!reinterpretShape.empty()) { shape = reinterpretShape; + } } - if (strides.empty()) + if (strides.empty()) { recordStaticSizes(reinterpret.getMixedStrides(), strides); + } return; } @@ -367,8 +410,9 @@ static void populateViewShapeAndStrides(Value value, } if (auto memrefType = dyn_cast(value.getType())) { - if (shape.empty()) + if (shape.empty()) { shape.assign(memrefType.getShape().begin(), memrefType.getShape().end()); + } if (strides.empty()) { int64_t offset = ShapedType::kDynamic; (void)mlir::pto::getPTOMemRefStridesAndOffset(memrefType, strides, @@ -379,8 +423,9 @@ static void populateViewShapeAndStrides(Value value, static std::optional getViewLayoutString(std::optional layout) { - if (!layout) + if (!layout) { return std::nullopt; + } return stringifyLayout(*layout).str(); } @@ -468,8 +513,9 @@ template static bool tryAppendPrecisionType( Operation *op, SmallVectorImpl> &attrs) { auto typed = dyn_cast(op); - if (!typed) + if (!typed) { return false; + } attrs.emplace_back("precisionType", getPrecisionTypeString(typed.getPrecisionType()).str()); return true; @@ -481,11 +527,13 @@ static bool tryAppendPrecisionType( static void appendOpContextAttrs( Operation *op, SmallVectorImpl> &attrs) { if (auto tcvt = dyn_cast(op)) { - if (auto roundMode = getTCvtRoundModeString(tcvt)) + if (auto roundMode = getTCvtRoundModeString(tcvt)) { attrs.emplace_back("round_mode", *roundMode); + } } - if (auto trandom = dyn_cast(op)) + if (auto trandom = dyn_cast(op)) { attrs.emplace_back("rounds", std::to_string(trandom.getRounds())); + } if (auto tcmp = dyn_cast(op)) { if (auto cmpModeAttr = tcmp.getCmpModeAttr()) attrs.emplace_back("cmp_mode", @@ -522,8 +570,9 @@ static void appendOpContextAttrs( } if (auto thistogram = dyn_cast(op)) { int byte = 1; - if (auto byteAttr = thistogram.getByteAttr()) + if (auto byteAttr = thistogram.getByteAttr()) { byte = byteAttr.getInt(); + } attrs.emplace_back("byte", std::to_string(byte)); } if (auto tscatter = dyn_cast(op)) { @@ -553,8 +602,9 @@ static std::string buildContextAttrsJson(Operation *operation) { std::string json = "{"; for (auto [index, attr] : llvm::enumerate(attrs)) { - if (index != 0) + if (index != 0) { json += ","; + } json += "\""; json += attr.first; json += "\":\""; @@ -585,8 +635,9 @@ static void appendTileOperandSpecJson(std::string &json, if (auto config = tileType.getConfigAttr()) { bLayout = config.getBLayout().getValue(); sLayout = config.getSLayout().getValue(); - if (config.getSFractalSize()) + if (config.getSFractalSize()) { fractalSize = config.getSFractalSize().getInt(); + } padValue = static_cast(config.getPad().getValue()); compactMode = static_cast(config.getCompactMode().getValue()); } @@ -611,8 +662,9 @@ static void appendViewOperandSpecJson(std::string &json, Value operand, SmallVector shape; SmallVector strides; populateViewShapeAndStrides(operand, shape, strides); - if (shape.empty()) + if (shape.empty()) { shape.assign(memrefType.getShape().begin(), memrefType.getShape().end()); + } appendJsonDimArray(json, shape); if (!strides.empty()) { json += ",\"strides\":"; @@ -636,8 +688,9 @@ static void appendViewOperandSpecJson(std::string &json, Value operand, SmallVector shape; SmallVector strides; populatePTOViewShapeAndStrides(operand, shape, strides); - if (shape.empty()) + if (shape.empty()) { shape.assign(viewType.getShape().begin(), viewType.getShape().end()); + } appendJsonDimArray(json, shape); if (!strides.empty()) { json += ",\"strides\":"; @@ -685,8 +738,9 @@ static std::optional buildOperandSpecsJson(Operation *operation) { std::string json = "["; for (auto [index, operand] : llvm::enumerate(operation->getOperands())) { - if (index != 0) + if (index != 0) { json += ","; + } Type type = operand.getType(); if (auto tileType = dyn_cast(type)) { @@ -764,8 +818,9 @@ getTargetArch(Operation *operation) { for (ModuleOp current = module; current; current = current->getParentOfType()) { - if (auto target = current->getAttrOfType("pto.target_arch")) + if (auto target = current->getAttrOfType("pto.target_arch")) { return target.getValue().str(); + } } operation->emitError( @@ -842,8 +897,9 @@ parseCandidateAttributes(Operation *operation, StringRef metadataJson) { llvm::sort(parsedCandidates, [](const CandidateMetadata &left, const CandidateMetadata &right) { - if (left.priority != right.priority) + if (left.priority != right.priority) { return left.priority > right.priority; + } return left.name < right.name; }); if (parsedCandidates.size() > 1 && @@ -891,11 +947,13 @@ struct InsertTemplateAttributesPass SmallVector tileOperations; module.walk([&](Operation *operation) { - if (pto::isTileLibExpandableOp(operation)) + if (pto::isTileLibExpandableOp(operation)) { tileOperations.push_back(operation); + } }); - if (tileOperations.empty()) + if (tileOperations.empty()) { return; + } std::shared_ptr tileLibService = pto::TileLibRuntime::getService(); if (!tileLibService) { @@ -907,8 +965,9 @@ struct InsertTemplateAttributesPass for (Operation *operation : tileOperations) { auto target = getTargetArch(operation); auto operandSpecs = buildOperandSpecsJson(operation); - if (!target || !operandSpecs) + if (!target || !operandSpecs) { return signalPassFailure(); + } pto::TileLibMaterializationRequest request; request.target = std::move(*target); request.op = operation->getName().getStringRef().str(); @@ -921,8 +980,9 @@ struct InsertTemplateAttributesPass } auto candidates = parseCandidateAttributes(operation, *metadata); - if (failed(candidates)) + if (failed(candidates)) { return signalPassFailure(); + } operation->setAttr(kCandidatesAttr, *candidates); } } diff --git a/lib/PTO/Transforms/LowerPTOToUBufOps.cpp b/lib/PTO/Transforms/LowerPTOToUBufOps.cpp index f9fe13866d..79c0cab5cc 100644 --- a/lib/PTO/Transforms/LowerPTOToUBufOps.cpp +++ b/lib/PTO/Transforms/LowerPTOToUBufOps.cpp @@ -41,7 +41,6 @@ namespace pto { } // namespace mlir namespace { - static constexpr int64_t kRepeatMax = 255; static constexpr int64_t kRepeatStrideMax = 255; static constexpr int64_t kSmallRptBinOp = 4; @@ -53,25 +52,31 @@ static constexpr unsigned kMaskLen = 64; //===----------------------------------------------------------------------===// static unsigned getElementSize(Type elemTy) { - if (elemTy.isF16() || elemTy.isBF16()) + if (elemTy.isF16() || elemTy.isBF16()) { return 2; - if (elemTy.isF32()) + } + if (elemTy.isF32()) { return 4; + } if (auto intTy = dyn_cast(elemTy)) { unsigned width = intTy.getWidth(); - if (width == 16 || width == 32) + if (width == 16 || width == 32) { return width / 8; + } } return 0; } static Type getStoredElemType(Type ty) { - if (auto tbTy = dyn_cast(ty)) + if (auto tbTy = dyn_cast(ty)) { return tbTy.getElementType(); - if (auto mrTy = dyn_cast(ty)) + } + if (auto mrTy = dyn_cast(ty)) { return mrTy.getElementType(); - if (auto ptrTy = dyn_cast(ty)) + } + if (auto ptrTy = dyn_cast(ty)) { return ptrTy.getElementType(); + } return Type(); } @@ -80,19 +85,22 @@ static std::optional isUBMemorySpaceImpl(Type ty) { if (auto tbTy = dyn_cast(ty)) { auto msAttr = dyn_cast_or_null(tbTy.getMemorySpace()); - if (!msAttr) + if (!msAttr) { return false; + } return msAttr.getAddressSpace() == pto::AddressSpace::VEC; } if (auto mrTy = dyn_cast(ty)) { auto msAttr = dyn_cast_or_null(mrTy.getMemorySpace()); - if (!msAttr) + if (!msAttr) { return false; + } return msAttr.getAddressSpace() == pto::AddressSpace::VEC; } - if (auto ptrTy = dyn_cast(ty)) + if (auto ptrTy = dyn_cast(ty)) { return ptrTy.getMemorySpace().getAddressSpace() == pto::AddressSpace::VEC; + } return std::nullopt; } @@ -104,8 +112,9 @@ static bool isUBMemorySpace(Type ty) { static bool isRowMajor(pto::TileBufType tbTy) { auto config = tbTy.getConfigAttr(); - if (!config) + if (!config) { return true; + } return config.getBLayout().getValue() != pto::BLayout::ColMajor; } @@ -145,8 +154,9 @@ using TileShapeMap = DenseMap; static std::optional extractTileShapeInfoFromValue( Value opDst, const TileShapeMap &tileShapes) { Type dstTy = opDst.getType(); - if (!isUBMemorySpace(dstTy)) + if (!isUBMemorySpace(dstTy)) { return std::nullopt; + } Type elemTy; ArrayRef shape; @@ -155,15 +165,17 @@ static std::optional extractTileShapeInfoFromValue( elemTy = tbTy.getElementType(); shape = tbTy.getShape(); validShape = tbTy.getValidShape(); - if (!isRowMajor(tbTy)) + if (!isRowMajor(tbTy)) { return std::nullopt; + } } else if (auto mrTy = dyn_cast(dstTy)) { elemTy = mrTy.getElementType(); shape = mrTy.getShape(); } else if (isa(dstTy)) { auto it = tileShapes.find(opDst); - if (it == tileShapes.end()) + if (it == tileShapes.end()) { return std::nullopt; + } elemTy = cast(dstTy).getElementType(); shape = llvm::ArrayRef(it->second.shape); validShape = llvm::ArrayRef(it->second.validShape); @@ -171,11 +183,13 @@ static std::optional extractTileShapeInfoFromValue( return std::nullopt; } unsigned elemSize = getElementSize(elemTy); - if (elemSize == 0) + if (elemSize == 0) { return std::nullopt; + } - if (shape.size() < 2) + if (shape.size() < 2) { return std::nullopt; + } int64_t rows = shape[0]; int64_t cols = shape[1]; @@ -186,8 +200,9 @@ static std::optional extractTileShapeInfoFromValue( validShape[1] != ShapedType::kDynamic) ? validShape[1] : cols; if (vRows == ShapedType::kDynamic || vCols == ShapedType::kDynamic || - rows == ShapedType::kDynamic || cols == ShapedType::kDynamic) + rows == ShapedType::kDynamic || cols == ShapedType::kDynamic) { return std::nullopt; + } TileShapeInfo info; info.vRows = vRows; @@ -219,15 +234,18 @@ struct LowerPTOToUBufOpsPass void runOnOperation() override { func::FuncOp func = getOperation(); - if (func.isExternal()) + if (func.isExternal()) { return; + } auto mod = func->getParentOfType(); - if (!mod) + if (!mod) { return; + } auto archAttr = mod->getAttrOfType("pto.target_arch"); if (!archAttr || - (archAttr.getValue() != "a2" && archAttr.getValue() != "a3")) + (archAttr.getValue() != "a2" && archAttr.getValue() != "a3")) { return; + } MLIRContext *ctx = &getContext(); OpBuilder builder(ctx); @@ -277,15 +295,18 @@ struct LowerPTOToUBufOpsPass SmallVector ops; func.walk([&](pto::TAddOp op) { ops.push_back(op); }); for (auto op : ops) { - if (!canLower(op, tileShapes)) + if (!canLower(op, tileShapes)) { continue; + } auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, src0Ptr, src1Ptr, ptrType] = lowerBinaryOpCommon( builder, ctx, op, op.getDst(), op.getSrc0(), op.getSrc1(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatch(op.getLoc(), builder, dstPtr, src0Ptr, src1Ptr, ptrType, *info); op.erase(); @@ -297,15 +318,18 @@ struct LowerPTOToUBufOpsPass SmallVector ops; func.walk([&](pto::TAddReluOp op) { ops.push_back(op); }); for (auto op : ops) { - if (!canLower(op, tileShapes)) + if (!canLower(op, tileShapes)) { continue; + } auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, src0Ptr, src1Ptr, ptrType] = lowerBinaryOpCommon( builder, ctx, op, op.getDst(), op.getSrc0(), op.getSrc1(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatch(op.getLoc(), builder, dstPtr, src0Ptr, src1Ptr, ptrType, *info); op.erase(); @@ -317,15 +341,18 @@ struct LowerPTOToUBufOpsPass SmallVector ops; func.walk([&](pto::TSubOp op) { ops.push_back(op); }); for (auto op : ops) { - if (!canLower(op, tileShapes)) + if (!canLower(op, tileShapes)) { continue; + } auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, src0Ptr, src1Ptr, ptrType] = lowerBinaryOpCommon( builder, ctx, op, op.getDst(), op.getSrc0(), op.getSrc1(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatch(op.getLoc(), builder, dstPtr, src0Ptr, src1Ptr, ptrType, *info); op.erase(); @@ -337,15 +364,18 @@ struct LowerPTOToUBufOpsPass SmallVector ops; func.walk([&](pto::TMulOp op) { ops.push_back(op); }); for (auto op : ops) { - if (!canLower(op, tileShapes)) + if (!canLower(op, tileShapes)) { continue; + } auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, src0Ptr, src1Ptr, ptrType] = lowerBinaryOpCommon( builder, ctx, op, op.getDst(), op.getSrc0(), op.getSrc1(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatch(op.getLoc(), builder, dstPtr, src0Ptr, src1Ptr, ptrType, *info); op.erase(); @@ -357,15 +387,18 @@ struct LowerPTOToUBufOpsPass SmallVector ops; func.walk([&](pto::TDivOp op) { ops.push_back(op); }); for (auto op : ops) { - if (!canLower(op, tileShapes)) + if (!canLower(op, tileShapes)) { continue; + } auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, src0Ptr, src1Ptr, ptrType] = lowerBinaryOpCommon( builder, ctx, op, op.getDst(), op.getSrc0(), op.getSrc1(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatch(op.getLoc(), builder, dstPtr, src0Ptr, src1Ptr, ptrType, *info); op.erase(); @@ -377,15 +410,18 @@ struct LowerPTOToUBufOpsPass SmallVector ops; func.walk([&](pto::TMaxOp op) { ops.push_back(op); }); for (auto op : ops) { - if (!canLower(op, tileShapes)) + if (!canLower(op, tileShapes)) { continue; + } auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, src0Ptr, src1Ptr, ptrType] = lowerBinaryOpCommon( builder, ctx, op, op.getDst(), op.getSrc0(), op.getSrc1(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatch(op.getLoc(), builder, dstPtr, src0Ptr, src1Ptr, ptrType, *info); op.erase(); @@ -397,15 +433,18 @@ struct LowerPTOToUBufOpsPass SmallVector ops; func.walk([&](pto::TMinOp op) { ops.push_back(op); }); for (auto op : ops) { - if (!canLower(op, tileShapes)) + if (!canLower(op, tileShapes)) { continue; + } auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, src0Ptr, src1Ptr, ptrType] = lowerBinaryOpCommon( builder, ctx, op, op.getDst(), op.getSrc0(), op.getSrc1(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatch(op.getLoc(), builder, dstPtr, src0Ptr, src1Ptr, ptrType, *info); op.erase(); @@ -417,15 +456,18 @@ struct LowerPTOToUBufOpsPass SmallVector ops; func.walk([&](pto::TAndOp op) { ops.push_back(op); }); for (auto op : ops) { - if (!canLower(op, tileShapes)) + if (!canLower(op, tileShapes)) { continue; + } auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, src0Ptr, src1Ptr, ptrType] = lowerBinaryOpCommon( builder, ctx, op, op.getDst(), op.getSrc0(), op.getSrc1(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatch(op.getLoc(), builder, dstPtr, src0Ptr, src1Ptr, ptrType, *info); op.erase(); @@ -437,15 +479,18 @@ struct LowerPTOToUBufOpsPass SmallVector ops; func.walk([&](pto::TOrOp op) { ops.push_back(op); }); for (auto op : ops) { - if (!canLower(op, tileShapes)) + if (!canLower(op, tileShapes)) { continue; + } auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, src0Ptr, src1Ptr, ptrType] = lowerBinaryOpCommon( builder, ctx, op, op.getDst(), op.getSrc0(), op.getSrc1(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatch(op.getLoc(), builder, dstPtr, src0Ptr, src1Ptr, ptrType, *info); op.erase(); @@ -459,13 +504,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TXorOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfoFromValue(op.getDst(), tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, src0Ptr, src1Ptr, tmpPtr, ptrType] = lowerXorOpCommon(builder, ctx, op, op.getDst(), op.getSrc0(), op.getSrc1(), op.getTmp(), tileShapes); - if (!dstPtr || !tmpPtr) + if (!dstPtr || !tmpPtr) { continue; + } auto pipeV = pto::PipeAttr::get(ctx, pto::PIPE::PIPE_V); // tmp = src0 | src1 dispatch(op.getLoc(), builder, tmpPtr, src0Ptr, src1Ptr, @@ -492,13 +539,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TNotOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfoFromValue(op.getDst(), tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatchUnary(op.getLoc(), builder, dstPtr, srcPtr, ptrType, *info); op.erase(); @@ -511,13 +560,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TAbsOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfoFromValue(op.getDst(), tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatchUnary(op.getLoc(), builder, dstPtr, srcPtr, ptrType, *info); op.erase(); @@ -530,13 +581,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TReluOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfoFromValue(op.getDst(), tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatchUnary(op.getLoc(), builder, dstPtr, srcPtr, ptrType, *info); op.erase(); @@ -549,13 +602,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TNegOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfoFromValue(op.getDst(), tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } Type elemTy = ptrType.getElementType(); Value minusOneScalar; if (elemTy.isF32() || elemTy.isF16()) { @@ -579,13 +634,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TRecipOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfoFromValue(op.getDst(), tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } Type elemTy = ptrType.getElementType(); Value oneScalar = builder.create( op.getLoc(), builder.getFloatAttr(elemTy, 1.0)); @@ -605,13 +662,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TExpOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfoFromValue(op.getDst(), tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatchUnary(op.getLoc(), builder, dstPtr, srcPtr, ptrType, *info); op.erase(); @@ -624,13 +683,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TLogOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfoFromValue(op.getDst(), tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatchUnary(op.getLoc(), builder, dstPtr, srcPtr, ptrType, *info); op.erase(); @@ -643,13 +704,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TSqrtOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfoFromValue(op.getDst(), tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatchUnary(op.getLoc(), builder, dstPtr, srcPtr, ptrType, *info); op.erase(); @@ -662,13 +725,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TRsqrtOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfoFromValue(op.getDst(), tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatchUnary(op.getLoc(), builder, dstPtr, srcPtr, ptrType, *info); op.erase(); @@ -681,13 +746,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TAddSOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } Value scalarI64 = convertScalarToI64(builder, op.getLoc(), op.getScalar()); dispatchShift(op.getLoc(), builder, dstPtr, srcPtr, scalarI64, ptrType, *info); @@ -701,13 +768,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TMulSOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc0(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } Value scalarI64 = convertScalarToI64(builder, op.getLoc(), op.getScalar()); dispatchShift(op.getLoc(), builder, dstPtr, srcPtr, scalarI64, ptrType, *info); @@ -721,13 +790,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TMaxSOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } Value scalarI64 = convertScalarToI64(builder, op.getLoc(), op.getScalar()); dispatchShift(op.getLoc(), builder, dstPtr, srcPtr, scalarI64, ptrType, *info); @@ -741,13 +812,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TMinSOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } Value scalarI64 = convertScalarToI64(builder, op.getLoc(), op.getScalar()); dispatchShift(op.getLoc(), builder, dstPtr, srcPtr, scalarI64, ptrType, *info); @@ -761,13 +834,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TShlSOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatchShift(op.getLoc(), builder, dstPtr, srcPtr, op.getScalar(), ptrType, *info); op.erase(); @@ -780,13 +855,15 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TShrSOp op) { ops.push_back(op); }); for (auto op : ops) { auto info = extractTileShapeInfo(op, tileShapes); - if (!info) + if (!info) { continue; + } auto [dstPtr, srcPtr, ptrType] = lowerShiftOpCommon(builder, ctx, op, op.getDst(), op.getSrc(), tileShapes); - if (!dstPtr) + if (!dstPtr) { continue; + } dispatchShift(op.getLoc(), builder, dstPtr, srcPtr, op.getScalar(), ptrType, *info); op.erase(); @@ -798,8 +875,9 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TLoadOp op) { tloadOps.push_back(op); }); for (auto op : tloadOps) { builder.setInsertionPoint(op); - if (succeeded(lowerTLoad(op, builder, tileShapes))) + if (succeeded(lowerTLoad(op, builder, tileShapes))) { op.erase(); + } } // ---- tstore → mte_ub_gm ---- @@ -807,20 +885,23 @@ struct LowerPTOToUBufOpsPass func.walk([&](pto::TStoreOp op) { tstoreOps.push_back(op); }); for (auto op : tstoreOps) { builder.setInsertionPoint(op); - if (succeeded(lowerTStore(op, builder, tileShapes))) + if (succeeded(lowerTStore(op, builder, tileShapes))) { op.erase(); + } } // ---- cleanup dead PTO ops ---- SmallVector toErase; func.walk([&](Operation *op) { if (isa(op)) + memref::SubViewOp, memref::ReinterpretCastOp, memref::CastOp>(op)) { toErase.push_back(op); + } }); for (auto *op : llvm::reverse(toErase)) { - if (op->use_empty()) + if (op->use_empty()) { op->erase(); + } } } @@ -879,8 +960,9 @@ struct LowerPTOToUBufOpsPass auto ptrType = getUBPtrType(ctx, elemTy); auto emitAddr = [&](Value tile) -> Value { - if (isa(tile.getType())) + if (isa(tile.getType())) { return tile; + } auto addrOp = builder.create(loc, ptrType, tile); return addrOp.getDst(); }; @@ -901,8 +983,9 @@ struct LowerPTOToUBufOpsPass auto ptrType = getUBPtrType(ctx, elemTy); auto emitAddr = [&](Value tile) -> Value { - if (isa(tile.getType())) + if (isa(tile.getType())) { return tile; + } auto addrOp = builder.create(loc, ptrType, tile); return addrOp.getDst(); }; @@ -921,8 +1004,9 @@ struct LowerPTOToUBufOpsPass Value asInt = builder.create(loc, intTy, scalar); return builder.create(loc, builder.getI64Type(), asInt); } - if (scalar.getType().isInteger(64)) + if (scalar.getType().isInteger(64)) { return scalar; + } return builder.create(loc, builder.getI64Type(), scalar); } @@ -935,8 +1019,9 @@ struct LowerPTOToUBufOpsPass auto ptrType = getUBPtrType(ctx, elemTy); auto emitAddr = [&](Value tile) -> Value { - if (isa(tile.getType())) + if (isa(tile.getType())) { return tile; + } auto addrOp = builder.create(loc, ptrType, tile); return addrOp.getDst(); }; @@ -972,9 +1057,10 @@ struct LowerPTOToUBufOpsPass auto emitShift = [&](Value d, Value s) { Value scalarI64 = scalar; - if (scalarI64.getType() != b.getI64Type()) + if (scalarI64.getType() != b.getI64Type()) { scalarI64 = b.create( loc, b.getI64Type(), scalar); + } b.create(loc, d, s, scalarI64, i64c1(loc, b), i64c1(loc, b), i64c1(loc, b), i64c8(loc, b), i64c8(loc, b)); @@ -1095,8 +1181,9 @@ struct LowerPTOToUBufOpsPass } auto emit = [&](Value rd) { Value scalarI64 = scalar; - if (scalarI64.getType() != b.getI64Type()) + if (scalarI64.getType() != b.getI64Type()) { scalarI64 = b.create(loc, b.getI64Type(), scalar); + } b.create(loc, rd, scalarI64, i64c1(loc, b), i64c1(loc, b), i64c1(loc, b), i64c8(loc, b), i64c0(loc, b)); @@ -1215,15 +1302,17 @@ struct LowerPTOToUBufOpsPass while (Operation *def = view.getDefiningOp()) { if (auto subview = dyn_cast(def)) { auto strides = subview.getStaticStrides(); - if (strides.empty() || strides.back() != 1) + if (strides.empty() || strides.back() != 1) { return false; + } view = subview.getSource(); continue; } if (auto reinterpret = dyn_cast(def)) { auto strides = reinterpret.getConstifiedMixedStrides(); - if (strides.empty()) + if (strides.empty()) { return false; + } auto stride = getConstantIntValue(strides.back()); return stride && *stride == 1; } @@ -1234,8 +1323,9 @@ struct LowerPTOToUBufOpsPass break; } auto memTy = dyn_cast(view.getType()); - if (!memTy) + if (!memTy) { return false; + } auto strides = memTy.getStridesAndOffset().first; return !strides.empty() && strides.back() == 1; } @@ -1244,8 +1334,9 @@ struct LowerPTOToUBufOpsPass auto pvOp = op.getSrc().getDefiningOp(); if (pvOp) { auto mtvOp = pvOp.getSource().getDefiningOp(); - if (!mtvOp) + if (!mtvOp) { return failure(); + } DmaViewInfo info; info.gmPtr = mtvOp.getPtr(); info.sizes.assign(pvOp.getSizes().begin(), pvOp.getSizes().end()); @@ -1266,8 +1357,9 @@ struct LowerPTOToUBufOpsPass auto pvOp = op.getDst().getDefiningOp(); if (pvOp) { auto mtvOp = pvOp.getSource().getDefiningOp(); - if (!mtvOp) + if (!mtvOp) { return failure(); + } DmaViewInfo info; info.gmPtr = mtvOp.getPtr(); info.sizes.assign(pvOp.getSizes().begin(), pvOp.getSizes().end()); @@ -1286,14 +1378,17 @@ struct LowerPTOToUBufOpsPass static FailureOr extractDmaMemRefViewInfo(Location loc, Value view, MLIRContext *ctx) { auto memTy = dyn_cast(view.getType()); - if (!memTy) + if (!memTy) { return failure(); + } auto msAttr = dyn_cast_or_null(memTy.getMemorySpace()); - if (!msAttr || msAttr.getAddressSpace() != pto::AddressSpace::GM) + if (!msAttr || msAttr.getAddressSpace() != pto::AddressSpace::GM) { return failure(); + } ArrayRef shape = memTy.getShape(); - if (shape.size() < 2) + if (shape.size() < 2) { return failure(); + } if (!hasUnitInnermostStride(view)) { emitError(loc) << "A2/A3 DMA lowering requires a unit innermost stride"; return failure(); @@ -1344,31 +1439,35 @@ struct LowerPTOToUBufOpsPass i < viewInfo.strides.size(); ++i) { APInt constOff; if (matchPattern(viewInfo.offsets[i], m_ConstantInt(&constOff)) && - constOff.isZero()) + constOff.isZero()) { continue; + } Value dimOff = b.create(loc, viewInfo.offsets[i], viewInfo.strides[i]).getResult(); totalOff = b.create(loc, totalOff, dimOff).getResult(); } - if (elemSize > 1) + if (elemSize > 1) { totalOff = b.create(loc, totalOff, idxc(elemSize, loc, b)).getResult(); + } return totalOff; } Value offsetGMPtrByBytes(Location loc, OpBuilder &b, Value gmPtr, Value byteOff) { APInt constOff; - if (matchPattern(byteOff, m_ConstantInt(&constOff)) && constOff.isZero()) + if (matchPattern(byteOff, m_ConstantInt(&constOff)) && constOff.isZero()) { return gmPtr; + } auto origPtrTy = cast(gmPtr.getType()); auto bytePtrTy = pto::PtrType::get(b.getContext(), b.getI8Type(), origPtrTy.getMemorySpace()); Value bytePtr = b.create(loc, bytePtrTy, gmPtr); Value offIdx = byteOff; - if (!offIdx.getType().isIndex()) + if (!offIdx.getType().isIndex()) { offIdx = b.create(loc, b.getIndexType(), byteOff) .getResult(); + } Value offsetBytePtr = b.create(loc, bytePtrTy, bytePtr, offIdx); return b.create(loc, origPtrTy, offsetBytePtr); @@ -1382,11 +1481,15 @@ struct LowerPTOToUBufOpsPass LogicalResult emitMteGmUb(Location loc, OpBuilder &b, Value gmPtr, Value ubPtr, const DmaViewInfo &viewInfo, Type elemTy, ArrayRef tileShape) { - if (tileShape.size() < 2) return failure(); + if (tileShape.size() < 2) { + return failure(); + } int64_t ubCols = tileShape[1]; unsigned elemSize = getElementSize(elemTy); unsigned nd = viewInfo.sizes.size(); - if (nd < 2) return failure(); + if (nd < 2) { + return failure(); + } Value nburstCount = viewInfo.sizes[nd - 2]; Value lenBurstElts = viewInfo.sizes[nd - 1]; @@ -1411,7 +1514,7 @@ struct LowerPTOToUBufOpsPass viewInfo.strides[i]).getResult(), i64c(elemSize, loc, b)).getResult(); Value innerElems = i64c1(loc, b); - for (int j = i + 1; j < (int)nd; ++j) { + for (int j = i + 1; j < static_cast(nd); ++j) { Value dimSize = b.create(loc, b.getI64Type(), viewInfo.sizes[j]).getResult(); innerElems = b.create(loc, innerElems, dimSize).getResult(); @@ -1428,11 +1531,15 @@ struct LowerPTOToUBufOpsPass LogicalResult emitMteUbGm(Location loc, OpBuilder &b, Value ubPtr, Value gmPtr, const DmaViewInfo &viewInfo, Type elemTy, ArrayRef tileShape) { - if (tileShape.size() < 2) return failure(); + if (tileShape.size() < 2) { + return failure(); + } int64_t ubCols = tileShape[1]; unsigned elemSize = getElementSize(elemTy); unsigned nd = viewInfo.sizes.size(); - if (nd < 2) return failure(); + if (nd < 2) { + return failure(); + } Value nburstCount = viewInfo.sizes[nd - 2]; Value lenBurstElts = viewInfo.sizes[nd - 1]; @@ -1453,7 +1560,7 @@ struct LowerPTOToUBufOpsPass Value count = b.create(loc, b.getI64Type(), viewInfo.sizes[i]).getResult(); Value innerElems = i64c1(loc, b); - for (int j = i + 1; j < (int)nd; ++j) { + for (int j = i + 1; j < static_cast(nd); ++j) { Value dimSize = b.create(loc, b.getI64Type(), viewInfo.sizes[j]).getResult(); innerElems = b.create(loc, innerElems, dimSize).getResult(); @@ -1475,16 +1582,26 @@ struct LowerPTOToUBufOpsPass const TileShapeMap &tileShapes) { Location loc = op.getLoc(); auto viewInfo = extractDmaViewInfo(op); - if (failed(viewInfo)) return failure(); + if (failed(viewInfo)) { + return failure(); + } Type dstType = op.getDst().getType(); - if (!isUBMemorySpace(dstType)) return failure(); + if (!isUBMemorySpace(dstType)) { + return failure(); + } Type elemTy = getStoredElemType(dstType); - if (!elemTy) return failure(); + if (!elemTy) { + return failure(); + } unsigned elemSize = getElementSize(elemTy); - if (elemSize == 0) return failure(); + if (elemSize == 0) { + return failure(); + } auto it = tileShapes.find(op.getDst()); - if (it == tileShapes.end()) return failure(); + if (it == tileShapes.end()) { + return failure(); + } Value byteOff = computeGMByteOffset(loc, b, *viewInfo, elemSize); Value gmPtr = offsetGMPtrByBytes(loc, b, viewInfo->gmPtr, byteOff); @@ -1496,17 +1613,27 @@ struct LowerPTOToUBufOpsPass const TileShapeMap &tileShapes) { Location loc = op.getLoc(); Type srcType = op.getSrc().getType(); - if (!isUBMemorySpace(srcType)) return failure(); + if (!isUBMemorySpace(srcType)) { + return failure(); + } Type elemTy = getStoredElemType(srcType); - if (!elemTy) return failure(); + if (!elemTy) { + return failure(); + } unsigned elemSize = getElementSize(elemTy); - if (elemSize == 0) return failure(); + if (elemSize == 0) { + return failure(); + } auto it = tileShapes.find(op.getSrc()); - if (it == tileShapes.end()) return failure(); + if (it == tileShapes.end()) { + return failure(); + } auto viewInfo = extractDmaViewInfo(op); - if (failed(viewInfo)) return failure(); + if (failed(viewInfo)) { + return failure(); + } Value byteOff = computeGMByteOffset(loc, b, *viewInfo, elemSize); Value gmPtr = offsetGMPtrByBytes(loc, b, viewInfo->gmPtr, byteOff); @@ -1538,10 +1665,11 @@ struct LowerPTOToUBufOpsPass int64_t totalV = vRows * vCols; int64_t totalRpts = (totalV + epr - 1) / epr; - if (totalRpts > kRepeatMax) + if (totalRpts > kRepeatMax) { modeCount1L(loc, b, dst, s0, s1, ptrTy, info); - else + } else { modeNorm1L(loc, b, dst, s0, s1, ptrTy, info); + } return; } @@ -1550,10 +1678,11 @@ struct LowerPTOToUBufOpsPass if (normColRepeat > 1 && vRows * normColRepeat < kSmallRptBinOp) { modeCount2L(loc, b, dst, s0, s1, ptrTy, info); } else if (vRows < normColRepeat + 1) { - if (vCols % epr > 0) + if (vCols % epr > 0) { modeCount2L(loc, b, dst, s0, s1, ptrTy, info); - else + } else { modeColVLAlign(loc, b, dst, s0, s1, ptrTy, info); + } } else { modeRowRpt(loc, b, dst, s0, s1, ptrTy, info); } @@ -1669,10 +1798,11 @@ struct LowerPTOToUBufOpsPass int64_t rs = rowStride / be; bool condRowRpt = (info.vRows <= kRepeatMax) && (rs <= kRepeatStrideMax); - if (condRowRpt) + if (condRowRpt) { rowRptFast(loc, b, dst, s0, s1, ptrTy, info, rs); - else + } else { rowRptChunked(loc, b, dst, s0, s1, ptrTy, info, rowStride, rs); + } } template @@ -1709,8 +1839,9 @@ struct LowerPTOToUBufOpsPass int64_t remainElem = info.vCols % epr; if (info.vRows > static_cast(epr)) { - if (rptPerLine > 0) + if (rptPerLine > 0) { headRows(loc, b, dst, s0, s1, ptrTy, info, rowStride, rptPerLine); + } if (remainElem > 0) { Value off = idxc(rptPerLine * epr, loc, b); tailRows(loc, b, addPtr(loc, b, dst, ptrTy, off), @@ -1802,20 +1933,22 @@ struct LowerPTOToUBufOpsPass idxc(numLoop, loc, b), idxc1(loc, b)); b.setInsertionPointToStart(forOp.getBody()); Value iv = forOp.getInductionVar(); - if (strideOver) + if (strideOver) { tailStrideOverChunk(loc, b, iv, dst, s0, s1, ptrTy, rowStride); - else + } else { tailStrideOkChunk(loc, b, iv, dst, s0, s1, ptrTy, rowStride, rs); + } b.setInsertionPointAfter(forOp); } if (remainAfterLoop > 0) { - if (strideOver) + if (strideOver) { tailStrideOverRemain(loc, b, dst, s0, s1, ptrTy, rowStride, numLoop, remainAfterLoop); - else + } else { tailStrideOkRemain(loc, b, dst, s0, s1, ptrTy, rowStride, rs, numLoop, remainAfterLoop); + } } fullMask(loc, b); @@ -1879,7 +2012,6 @@ struct LowerPTOToUBufOpsPass i64c(remain, loc, b), i64c(rs, loc, b)); } }; - } // namespace namespace mlir { diff --git a/lib/PTO/Transforms/PTOA5NormalizeTMovPass.cpp b/lib/PTO/Transforms/PTOA5NormalizeTMovPass.cpp index 852273ed22..8b5e47aabf 100644 --- a/lib/PTO/Transforms/PTOA5NormalizeTMovPass.cpp +++ b/lib/PTO/Transforms/PTOA5NormalizeTMovPass.cpp @@ -128,11 +128,13 @@ static pto::TileBufConfigAttr buildRowMajorConfig(MLIRContext *ctx, static FailureOr buildRowMajorReinterpretType(MLIRContext *ctx, pto::TileBufType srcType) { ArrayRef shape = srcType.getShape(); - if (shape.size() != kTileRank2D) + if (shape.size() != kTileRank2D) { return failure(); + } if (shape[kFirstTileDim] == ShapedType::kDynamic || - shape[kSecondTileDim] == ShapedType::kDynamic) + shape[kSecondTileDim] == ShapedType::kDynamic) { return failure(); + } SmallVector swappedShape{shape[kSecondTileDim], shape[kFirstTileDim]}; @@ -160,8 +162,9 @@ buildRowMajorReinterpretType(MLIRContext *ctx, pto::TileBufType srcType) { static void setSwappedDynamicValidShapeIfNeeded( IRRewriter &rewriter, Location loc, Value sourceTile, Value reshapedTile, pto::TileBufType reshapedType) { - if (!reshapedType.hasDynamicValid()) + if (!reshapedType.hasDynamicValid()) { return; + } auto validShape = rewriter.create(loc, sourceTile); rewriter.create( @@ -179,15 +182,17 @@ struct PTOA5NormalizeTMovPass func.walk([&](pto::TGetScaleAddrOp op) { scaleAddrOps.push_back(op); }); for (pto::TGetScaleAddrOp op : scaleAddrOps) { auto matchingTMov = findMatchingScaleTileTMov(op); - if (!matchingTMov) + if (!matchingTMov) { continue; + } op->moveBefore(matchingTMov); } SmallVector riskyOps; func.walk([&](pto::TMovOp op) { - if (isA5RiskyVecVecColMajorTMov(op)) + if (isA5RiskyVecVecColMajorTMov(op)) { riskyOps.push_back(op); + } }); IRRewriter rewriter(func.getContext()); @@ -238,16 +243,18 @@ struct PTOA5NormalizeTMovPass bool hasResidualRisk = false; func.walk([&](pto::TMovOp op) { - if (!isA5RiskyVecVecColMajorTMov(op)) + if (!isA5RiskyVecVecColMajorTMov(op)) { return WalkResult::advance(); + } op.emitOpError( "A5 vec->vec TMOV on col_major/none_box tile is unsupported; " "expected normalization to row_major via pto.treshape"); hasResidualRisk = true; return WalkResult::interrupt(); }); - if (hasResidualRisk) + if (hasResidualRisk) { signalPassFailure(); + } } }; diff --git a/lib/PTO/Transforms/PTOAssignDefaultFrontendPipeIdPass.cpp b/lib/PTO/Transforms/PTOAssignDefaultFrontendPipeIdPass.cpp index 5f4f2bd115..187174ee7d 100644 --- a/lib/PTO/Transforms/PTOAssignDefaultFrontendPipeIdPass.cpp +++ b/lib/PTO/Transforms/PTOAssignDefaultFrontendPipeIdPass.cpp @@ -25,8 +25,9 @@ namespace { template static void assignDefaultIdIfMissing(OpT op, IntegerAttr zeroAttr) { - if (!op.getIdAttr()) + if (!op.getIdAttr()) { op.setIdAttr(zeroAttr); + } } struct PTOAssignDefaultFrontendPipeIdPass diff --git a/lib/PTO/Transforms/PTOCanonicalizeIR.cpp b/lib/PTO/Transforms/PTOCanonicalizeIR.cpp index 3f95ce38e4..d39106b896 100644 --- a/lib/PTO/Transforms/PTOCanonicalizeIR.cpp +++ b/lib/PTO/Transforms/PTOCanonicalizeIR.cpp @@ -129,8 +129,9 @@ buildCanonicalStrides(MakeTensorViewOp op, IRRewriter &rewriter) { Value leadingStride = rewriter.create( loc, op.getShape().front(), op.getStrides().front()); - for (unsigned i = 0; i < shift; ++i) + for (unsigned i = 0; i < shift; ++i) { result[i] = leadingStride; + } return result; } @@ -177,8 +178,9 @@ static Type canonicalViewType(Type type) { static bool canonicalizeValueType(Value value) { Type oldType = value.getType(); Type newType = canonicalViewType(oldType); - if (newType == oldType) + if (newType == oldType) { return false; + } value.setType(newType); return true; } @@ -277,8 +279,9 @@ static void canonicalizeFunctionType(func::FuncOp func) { results.push_back(newType); } - if (changed) + if (changed) { func.setFunctionType(FunctionType::get(func.getContext(), inputs, results)); + } } static void canonicalizeValueTypes(func::FuncOp func) { @@ -362,8 +365,9 @@ struct PTOCanonicalizeIRPass return; } } - for (auto [op, dimIndex, rank] : dimIndexOps) + for (auto [op, dimIndex, rank] : dimIndexOps) { rewriteTensorViewDimOperand(op, dimIndex, rank, rewriter); + } canonicalizeValueTypes(func); for (PartitionViewOp op : partitionViews) { if (failed(rewritePartitionView(op, rewriter))) { diff --git a/lib/PTO/Transforms/PTOInferVPTOVecScope.cpp b/lib/PTO/Transforms/PTOInferVPTOVecScope.cpp index 834fcd6ec8..d0b2975768 100644 --- a/lib/PTO/Transforms/PTOInferVPTOVecScope.cpp +++ b/lib/PTO/Transforms/PTOInferVPTOVecScope.cpp @@ -102,19 +102,22 @@ static bool isVectorScopeBoundaryOperation(Operation *op) { static bool hasVecScopeTypedOperandOrResult(Operation *op) { for (Type type : op->getOperandTypes()) { - if (isVecScopeType(type)) + if (isVecScopeType(type)) { return true; + } } for (Type type : op->getResultTypes()) { - if (isVecScopeType(type)) + if (isVecScopeType(type)) { return true; + } } return false; } static bool requiresVectorScope(Operation *op) { - if (!isPTOOperation(op)) + if (!isPTOOperation(op)) { return false; + } return hasVecScopeTypedOperandOrResult(op) || isa(op); @@ -125,24 +128,31 @@ static bool isAtomicControlFlowCandidate(Operation *op) { } static bool isSafeScalarOperation(Operation *op) { - if (op->getNumRegions() != 0) + if (op->getNumRegions() != 0) { return false; - if (op->hasTrait()) + } + if (op->hasTrait()) { return false; - if (isa(op)) + } + if (isa(op)) { return false; - if (isPTOOperation(op) && !isMemoryEffectFree(op)) + } + if (isPTOOperation(op) && !isMemoryEffectFree(op)) { return false; + } return isMemoryEffectFree(op); } static bool isRematerializableVecScopeProducer(Operation *op) { - if (!op || op->getNumRegions() != 0) + if (!op || op->getNumRegions() != 0) { return false; - if (op->hasTrait()) + } + if (op->hasTrait()) { return false; - if (isa(op)) + } + if (isa(op)) { return false; + } return isMemoryEffectFree(op); } @@ -150,8 +160,9 @@ static void summarizeNestedRegionForAtomicCluster( Region ®ion, NestedRegionSummary &summary) { for (Block &block : region) { for (Operation &op : block) { - if (op.hasTrait()) + if (op.hasTrait()) { continue; + } switch (classifyOperationForInference(&op)) { case VPTOInferenceOpClass::Vector: @@ -168,41 +179,52 @@ static void summarizeNestedRegionForAtomicCluster( } static bool canTreatAsAtomicControlFlow(Operation *op) { - if (!isAtomicControlFlowCandidate(op)) + if (!isAtomicControlFlowCandidate(op)) { return false; + } NestedRegionSummary summary; for (Region ®ion : op->getRegions()) { summarizeNestedRegionForAtomicCluster(region, summary); - if (summary.hasBoundaryOperation) + if (summary.hasBoundaryOperation) { return false; + } } return summary.hasVectorOperation; } static VPTOInferenceOpClass classifyOperationForInference(Operation *op) { - if (!op) + if (!op) { return VPTOInferenceOpClass::Boundary; + } - if (isExplicitVectorScopeCarrier(op)) + if (isExplicitVectorScopeCarrier(op)) { return VPTOInferenceOpClass::Boundary; - if (op->hasTrait()) + } + if (op->hasTrait()) { return VPTOInferenceOpClass::Boundary; - if (isa(op)) + } + if (isa(op)) { return VPTOInferenceOpClass::Boundary; - if (isVectorScopeBoundaryOperation(op)) + } + if (isVectorScopeBoundaryOperation(op)) { return VPTOInferenceOpClass::Boundary; - if (isForbiddenInsideInferredVectorScope(op)) + } + if (isForbiddenInsideInferredVectorScope(op)) { return VPTOInferenceOpClass::Boundary; + } - if (requiresVectorScope(op)) + if (requiresVectorScope(op)) { return VPTOInferenceOpClass::Vector; + } - if (canTreatAsAtomicControlFlow(op)) + if (canTreatAsAtomicControlFlow(op)) { return VPTOInferenceOpClass::Vector; + } - if (isSafeScalarOperation(op)) + if (isSafeScalarOperation(op)) { return VPTOInferenceOpClass::SafeScalar; + } return VPTOInferenceOpClass::Boundary; } @@ -216,8 +238,9 @@ static bool hasVectorOperation(ArrayRef ops) { static bool isUserInsideCluster(Operation *user, const llvm::SmallPtrSetImpl &ops) { for (Operation *cur = user; cur; cur = cur->getParentOp()) { - if (ops.contains(cur)) + if (ops.contains(cur)) { return true; + } } return false; } @@ -225,8 +248,9 @@ static bool isUserInsideCluster(Operation *user, static bool anyUserIsMoved(Value result, const llvm::SmallPtrSetImpl &movedOps) { for (Operation *user : result.getUsers()) { - if (isUserInsideCluster(user, movedOps)) + if (isUserInsideCluster(user, movedOps)) { return true; + } } return false; } @@ -235,8 +259,9 @@ static llvm::SmallPtrSet computeMovedOpsForResultlessScope(ArrayRef ops) { llvm::SmallPtrSet movedOps; for (Operation *op : ops) { - if (classifyOperationForInference(op) == VPTOInferenceOpClass::Vector) + if (classifyOperationForInference(op) == VPTOInferenceOpClass::Vector) { movedOps.insert(op); + } } bool changed = true; @@ -261,8 +286,9 @@ computeMovedOpsForResultlessScope(ArrayRef ops) { break; } } - if (!allUsersMoved) + if (!allUsersMoved) { break; + } } if (hasMovedUser && allUsersMoved) { @@ -276,8 +302,9 @@ computeMovedOpsForResultlessScope(ArrayRef ops) { static Operation *getAncestorInBlock(Operation *op, Block &block) { for (Operation *cur = op; cur; cur = cur->getParentOp()) { - if (cur->getBlock() == &block) + if (cur->getBlock() == &block) { return cur; + } } return nullptr; } @@ -288,33 +315,39 @@ cloneVecScopeProducerForUse( SegmentRematCache &cache, MLIRContext *context, llvm::DenseMap &clones) { auto result = dyn_cast(value); - if (!result) + if (!result) { return failure(); + } if (auto cacheIt = cache.find(value); cacheIt != cache.end()) { auto anchorIt = cacheIt->second.find(logicalScopeAnchor); - if (anchorIt != cacheIt->second.end()) + if (anchorIt != cacheIt->second.end()) { return anchorIt->second.getDefiningOp(); + } } Operation *producer = result.getOwner(); auto existing = clones.find(producer); - if (existing != clones.end()) + if (existing != clones.end()) { return existing->second; + } - if (!isRematerializableVecScopeProducer(producer)) + if (!isRematerializableVecScopeProducer(producer)) { return failure(); + } IRMapping mapping; for (Value operand : producer->getOperands()) { - if (!isVecScopeType(operand.getType())) + if (!isVecScopeType(operand.getType())) { continue; + } FailureOr clonedOperandProducer = cloneVecScopeProducerForUse(operand, user, logicalScopeAnchor, cache, context, clones); - if (failed(clonedOperandProducer)) + if (failed(clonedOperandProducer)) { return failure(); + } auto operandResult = cast(operand); mapping.map(operand, (*clonedOperandProducer) @@ -339,8 +372,9 @@ static void collectGreedyLogicalScopePlans( for (size_t end = ops.size(); end > begin; --end) { ArrayRef candidate = ops.slice(begin, end - begin); - if (!hasVectorOperation(candidate)) + if (!hasVectorOperation(candidate)) { continue; + } ResultlessScopePlan plan; EscapingMovedValue candidateEscapingValue; @@ -370,11 +404,13 @@ static void assignLogicalScopeAnchorsForCluster( llvm::DenseMap scopeAnchorByMovedOp; for (const LogicalScopePlan &plan : plans) { - if (plan.plan.moveOps.empty()) + if (plan.plan.moveOps.empty()) { continue; + } Operation *scopeAnchor = plan.plan.moveOps.front(); - for (Operation *movedOp : plan.plan.moveOps) + for (Operation *movedOp : plan.plan.moveOps) { scopeAnchorByMovedOp[movedOp] = scopeAnchor; + } } Operation *currentNonScopeAnchor = nullptr; @@ -386,8 +422,9 @@ static void assignLogicalScopeAnchorsForCluster( continue; } - if (!currentNonScopeAnchor) + if (!currentNonScopeAnchor) { currentNonScopeAnchor = op; + } logicalScopeAnchors[op] = currentNonScopeAnchor; } } @@ -398,8 +435,9 @@ computeLogicalScopeAnchors(Block &block) { SmallVector pending; auto flush = [&]() { - if (pending.empty()) + if (pending.empty()) { return; + } assignLogicalScopeAnchorsForCluster(pending, logicalScopeAnchors); pending.clear(); }; @@ -423,8 +461,9 @@ static LogicalResult rematerializeEscapingValueForUserSegments( Value value, const llvm::SmallPtrSetImpl &movedOps, Block &block, SegmentRematCache &cache, MLIRContext *context) { auto result = dyn_cast(value); - if (!result) + if (!result) { return failure(); + } llvm::DenseMap logicalScopeAnchors = computeLogicalScopeAnchors(block); @@ -432,22 +471,26 @@ static LogicalResult rematerializeEscapingValueForUserSegments( for (OpOperand &use : result.getUses()) { Operation *user = use.getOwner(); - if (isUserInsideCluster(user, movedOps)) + if (isUserInsideCluster(user, movedOps)) { continue; + } Operation *ancestor = getAncestorInBlock(user, block); - if (!ancestor) + if (!ancestor) { return failure(); + } auto anchorIt = logicalScopeAnchors.find(ancestor); - if (anchorIt == logicalScopeAnchors.end()) + if (anchorIt == logicalScopeAnchors.end()) { return failure(); + } usesBySegment[anchorIt->second].push_back(&use); } - if (usesBySegment.empty()) + if (usesBySegment.empty()) { return failure(); + } for (auto &entry : usesBySegment) { Operation *logicalScopeAnchor = entry.first; @@ -455,28 +498,32 @@ static LogicalResult rematerializeEscapingValueForUserSegments( Value replacement; if (auto cacheIt = cache.find(value); cacheIt != cache.end()) { auto anchorIt = cacheIt->second.find(logicalScopeAnchor); - if (anchorIt != cacheIt->second.end()) + if (anchorIt != cacheIt->second.end()) { replacement = anchorIt->second; + } } if (!replacement) { - if (!logicalScopeAnchor) + if (!logicalScopeAnchor) { return failure(); + } llvm::DenseMap clones; FailureOr clonedProducer = cloneVecScopeProducerForUse(value, logicalScopeAnchor, logicalScopeAnchor, cache, context, clones); - if (failed(clonedProducer)) + if (failed(clonedProducer)) { return failure(); + } replacement = (*clonedProducer)->getResult(result.getResultNumber()); cache[value][logicalScopeAnchor] = replacement; } - for (OpOperand *use : uses) + for (OpOperand *use : uses) { use->set(replacement); + } } return success(); @@ -488,8 +535,9 @@ static bool findEscapingMovedResult( for (Operation *op : movedOps) { for (Value result : op->getResults()) { for (Operation *user : result.getUsers()) { - if (isUserInsideCluster(user, movedOps)) + if (isUserInsideCluster(user, movedOps)) { continue; + } escapingValue.value = result; escapingValue.producer = op; @@ -505,18 +553,21 @@ static bool findEscapingMovedResult( static LogicalResult emitEscapingVectorScopeValueError(const EscapingMovedValue &escapingValue) { Operation *producer = escapingValue.producer; - if (!producer) + if (!producer) { return failure(); + } InFlightDiagnostic diag = producer->emitOpError() << "cannot infer resultless pto.vecscope because " "VPTO vector-scope data cannot have external " "users"; - if (escapingValue.value) + if (escapingValue.value) { diag << "; escaping value type is " << escapingValue.value.getType(); - if (escapingValue.user) + } + if (escapingValue.user) { diag.attachNote(escapingValue.user->getLoc()) << "external user is here"; + } return failure(); } @@ -527,16 +578,19 @@ emitEscapingVectorScopeValueError(const EscapingMovedValue &escapingValue) { static LogicalResult buildResultlessScopePlan(ArrayRef ops, ResultlessScopePlan &plan, EscapingMovedValue &escapingValue) { - if (ops.empty() || !hasVectorOperation(ops)) + if (ops.empty() || !hasVectorOperation(ops)) { return failure(); + } llvm::SmallPtrSet movedOps = computeMovedOpsForResultlessScope(ops); - if (movedOps.empty()) + if (movedOps.empty()) { return failure(); + } - if (findEscapingMovedResult(movedOps, escapingValue)) + if (findEscapingMovedResult(movedOps, escapingValue)) { return failure(); + } llvm::SmallPtrSet hoistedOps; for (Operation *op : ops) { @@ -569,8 +623,9 @@ buildResultlessScopePlan(ArrayRef ops, ResultlessScopePlan &plan, break; } } - if (feedsHoistedOp) + if (feedsHoistedOp) { break; + } } if (feedsHoistedOp) { @@ -583,17 +638,20 @@ buildResultlessScopePlan(ArrayRef ops, ResultlessScopePlan &plan, plan.hoistOps.clear(); plan.moveOps.clear(); for (Operation *op : ops) { - if (hoistedOps.contains(op)) + if (hoistedOps.contains(op)) { plan.hoistOps.push_back(op); - if (movedOps.contains(op)) + } + if (movedOps.contains(op)) { plan.moveOps.push_back(op); + } } return success(); } static void wrapCluster(const ResultlessScopePlan &plan, MLIRContext *context) { - if (plan.moveOps.empty()) + if (plan.moveOps.empty()) { return; + } Operation *first = plan.moveOps.front(); Block *parentBlock = first->getBlock(); @@ -604,8 +662,9 @@ static void wrapCluster(const ResultlessScopePlan &plan, MLIRContext *context) { scope.getBody().push_back(new Block()); for (Operation *op : plan.hoistOps) { - if (op->getBlock() == parentBlock && scope->isBeforeInBlock(op)) + if (op->getBlock() == parentBlock && scope->isBeforeInBlock(op)) { op->moveBefore(scope); + } } Block &scopeBody = scope.getBody().front(); @@ -626,8 +685,9 @@ static LogicalResult wrapGreedySubclusters(ArrayRef ops, for (size_t end = ops.size(); end > begin; --end) { ArrayRef candidate = ops.slice(begin, end - begin); - if (!hasVectorOperation(candidate)) + if (!hasVectorOperation(candidate)) { continue; + } // Prefer the largest suffix-preserving candidate that actually needs a // vecscope and can be moved into today's resultless pto.vecscope form. @@ -673,8 +733,9 @@ static FailureOr fixOneEscapingSubcluster(ArrayRef ops, for (size_t end = ops.size(); end > begin; --end) { ArrayRef candidate = ops.slice(begin, end - begin); - if (!hasVectorOperation(candidate)) + if (!hasVectorOperation(candidate)) { continue; + } ResultlessScopePlan ignoredPlan; EscapingMovedValue candidateEscapingValue; @@ -702,8 +763,9 @@ static FailureOr fixOneEscapingSubcluster(ArrayRef ops, llvm::SmallPtrSet movedOps = computeMovedOpsForResultlessScope(escapingCandidate); Block *block = ops.front()->getBlock(); - if (!block) + if (!block) { return false; + } if (succeeded(rematerializeEscapingValueForUserSegments( escapingValue.value, movedOps, *block, cache, context))) @@ -728,8 +790,9 @@ static LogicalResult repairEscapingSubclusters(Block &block, bool changedInIteration = false; SmallVector pending; SmallVector ops; - for (Operation &op : block) + for (Operation &op : block) { ops.push_back(&op); + } auto flush = [&]() -> FailureOr { FailureOr changed = @@ -746,45 +809,53 @@ static LogicalResult repairEscapingSubclusters(Block &block, break; case VPTOInferenceOpClass::Boundary: { FailureOr changed = flush(); - if (failed(changed)) + if (failed(changed)) { return failure(); + } changedInIteration |= *changed; break; } } - if (changedInIteration) + if (changedInIteration) { break; + } } - if (changedInIteration) + if (changedInIteration) { continue; + } FailureOr changed = flush(); - if (failed(changed)) + if (failed(changed)) { return failure(); - if (!*changed) + } + if (!*changed) { return success(); + } } return failure(); } static LogicalResult inferVecScopesInBlock(Block &block, MLIRContext *context) { - if (failed(repairEscapingSubclusters(block, context))) + if (failed(repairEscapingSubclusters(block, context))) { return failure(); + } SmallVector pending; auto flush = [&]() -> LogicalResult { - if (failed(wrapGreedySubclusters(pending, context))) + if (failed(wrapGreedySubclusters(pending, context))) { return failure(); + } pending.clear(); return success(); }; SmallVector ops; - for (Operation &op : block) + for (Operation &op : block) { ops.push_back(&op); + } for (Operation *op : ops) { switch (classifyOperationForInference(op)) { @@ -793,24 +864,29 @@ static LogicalResult inferVecScopesInBlock(Block &block, MLIRContext *context) { pending.push_back(op); continue; case VPTOInferenceOpClass::Boundary: - if (failed(flush())) + if (failed(flush())) { return failure(); + } continue; } } - if (failed(flush())) + if (failed(flush())) { return failure(); + } SmallVector remainingOps; - for (Operation &op : block) + for (Operation &op : block) { remainingOps.push_back(&op); + } for (Operation *op : remainingOps) { - if (isExplicitVectorScopeCarrier(op)) + if (isExplicitVectorScopeCarrier(op)) { continue; + } for (Region &nested : op->getRegions()) { - if (failed(inferVecScopesInRegion(nested, context))) + if (failed(inferVecScopesInRegion(nested, context))) { return failure(); + } } } return success(); @@ -819,8 +895,9 @@ static LogicalResult inferVecScopesInBlock(Block &block, MLIRContext *context) { static LogicalResult inferVecScopesInRegion(Region ®ion, MLIRContext *context) { for (Block &block : region) { - if (failed(inferVecScopesInBlock(block, context))) + if (failed(inferVecScopesInBlock(block, context))) { return failure(); + } } return success(); } @@ -830,8 +907,9 @@ struct PTOInferVPTOVecScopePass PTOInferVPTOVecScopePass> { void runOnOperation() override { func::FuncOp func = getOperation(); - if (failed(inferVecScopesInRegion(func.getBody(), &getContext()))) + if (failed(inferVecScopesInRegion(func.getBody(), &getContext()))) { signalPassFailure(); + } } }; diff --git a/lib/PTO/Transforms/PTOInferValidatePipeInitPass.cpp b/lib/PTO/Transforms/PTOInferValidatePipeInitPass.cpp index d5878a59c0..d9dc5f99c9 100644 --- a/lib/PTO/Transforms/PTOInferValidatePipeInitPass.cpp +++ b/lib/PTO/Transforms/PTOInferValidatePipeInitPass.cpp @@ -77,8 +77,9 @@ template static Value getLocalAddrOperand(InitOpT op) { template static std::optional getNoSplitAttr(InitOpT op) { - if (auto attr = op.getNosplitAttr()) + if (auto attr = op.getNosplitAttr()) { return attr.getValue(); + } return std::nullopt; } @@ -105,19 +106,23 @@ static PipeSplitUsage classifyPipeUsage(Value pipe) { continue; } - if (split == 0) + if (split == 0) { sawNoSplit = true; - else + } else { sawSplit = true; + } - if (sawNoSplit && sawSplit) + if (sawNoSplit && sawSplit) { return PipeSplitUsage::Mixed; + } } - if (sawNoSplit) + if (sawNoSplit) { return PipeSplitUsage::NoSplitOnly; - if (sawSplit) + } + if (sawSplit) { return PipeSplitUsage::SplitOnly; + } return PipeSplitUsage::Unknown; } @@ -141,10 +146,11 @@ static std::string getFuncSymbol(func::FuncOp funcOp) { static PipePeerKey getGlobalTensorPipeKey(Operation *op, int8_t dirMask) { std::string id = "unknown"; - if (auto idAttr = op->getAttrOfType(kFrontendPipeIdAttrName)) + if (auto idAttr = op->getAttrOfType(kFrontendPipeIdAttrName)) { id = std::to_string(idAttr.getInt()); - else + } else { id = std::to_string(reinterpret_cast(op)); + } return PipePeerKey{"__pto_globaltensor_pipe", "id_" + id, dirMask}; } @@ -176,8 +182,9 @@ resolveNoSplitComponent(ArrayRef component, OpBuilder &builder) "same logical pipe"); } - if (!info->explicitNoSplit) + if (!info->explicitNoSplit) { continue; + } if (explicitNoSplit && *explicitNoSplit != *info->explicitNoSplit) { return info->op->emitOpError( "conflicting explicit 'nosplit' across peer pipe init ops"); @@ -187,8 +194,9 @@ resolveNoSplitComponent(ArrayRef component, OpBuilder &builder) for (PipeInitInfo *info : component) { auto usageNoSplit = getUsageNoSplit(info->usage); - if (!usageNoSplit) + if (!usageNoSplit) { continue; + } if (inferredNoSplit && *inferredNoSplit != *usageNoSplit) { return info->op->emitOpError( "conflicting pipe split usage across peer pipe init ops"); @@ -198,8 +206,9 @@ resolveNoSplitComponent(ArrayRef component, OpBuilder &builder) if (explicitNoSplit && inferredNoSplit && *explicitNoSplit != *inferredNoSplit) { for (PipeInitInfo *info : component) { - if (!info->explicitNoSplit || *info->explicitNoSplit == *inferredNoSplit) + if (!info->explicitNoSplit || *info->explicitNoSplit == *inferredNoSplit) { continue; + } if (*info->explicitNoSplit) { return info->op->emitOpError( "explicit 'nosplit = true' conflicts with downstream users that " @@ -216,14 +225,16 @@ resolveNoSplitComponent(ArrayRef component, OpBuilder &builder) auto noSplitAttr = builder.getBoolAttr(finalNoSplit); for (PipeInitInfo *info : component) { if (auto initOp = dyn_cast(info->op)) { - if (!initOp.getNosplitAttr()) + if (!initOp.getNosplitAttr()) { setNoSplitAttr(initOp, noSplitAttr); + } continue; } auto initOp = cast(info->op); - if (!initOp.getNosplitAttr()) + if (!initOp.getNosplitAttr()) { setNoSplitAttr(initOp, noSplitAttr); + } } return success(); @@ -248,11 +259,13 @@ struct PTOInferValidatePipeInitPass adjacency[info.op]; auto recordAddr = [&](Value addr, int8_t effectiveDirMask) { - if (!addr) + if (!addr) { return; + } auto key = getPipePeerKey(addr, info.funcOp); - if (!key) + if (!key) { return; + } key->dirMask = effectiveDirMask; keyedInits[*key].push_back(info.op); }; @@ -275,8 +288,9 @@ struct PTOInferValidatePipeInitPass if (info.dirMask == kBidirectionalDirMask) { recordAddr(getLocalAddrOperand(initOp), kC2VDirMask); - if (Value peerAddr = initOp.getPeerLocalAddr()) + if (Value peerAddr = initOp.getPeerLocalAddr()) { recordAddr(peerAddr, kV2CDirMask); + } return; } @@ -289,11 +303,13 @@ struct PTOInferValidatePipeInitPass for (const auto &it : keyedInits) { SmallVector uniqueOps; for (Operation *op : it.second) { - if (std::find(uniqueOps.begin(), uniqueOps.end(), op) == uniqueOps.end()) + if (std::find(uniqueOps.begin(), uniqueOps.end(), op) == uniqueOps.end()) { uniqueOps.push_back(op); + } } - if (uniqueOps.size() < kMinPeerPipeInitCount) + if (uniqueOps.size() < kMinPeerPipeInitCount) { continue; + } for (size_t i = 0; i < uniqueOps.size(); ++i) { for (size_t j = i + 1; j < uniqueOps.size(); ++j) { @@ -304,14 +320,16 @@ struct PTOInferValidatePipeInitPass } llvm::DenseMap infoByOp; - for (PipeInitInfo &info : initInfos) + for (PipeInitInfo &info : initInfos) { infoByOp[info.op] = &info; + } OpBuilder builder(moduleOp.getContext()); llvm::SmallPtrSet visited; for (PipeInitInfo &rootInfo : initInfos) { - if (!visited.insert(rootInfo.op).second) + if (!visited.insert(rootInfo.op).second) { continue; + } SmallVector stack{rootInfo.op}; SmallVector component; @@ -319,8 +337,9 @@ struct PTOInferValidatePipeInitPass Operation *current = stack.pop_back_val(); component.push_back(infoByOp[current]); for (Operation *neighbor : adjacency[current]) { - if (visited.insert(neighbor).second) + if (visited.insert(neighbor).second) { stack.push_back(neighbor); + } } } diff --git a/lib/PTO/Transforms/PTOInstantiateAndInlineOpLib.cpp b/lib/PTO/Transforms/PTOInstantiateAndInlineOpLib.cpp index 85a8d6f2a3..25cc2c2ba1 100644 --- a/lib/PTO/Transforms/PTOInstantiateAndInlineOpLib.cpp +++ b/lib/PTO/Transforms/PTOInstantiateAndInlineOpLib.cpp @@ -87,11 +87,13 @@ static Operation *cloneOpForInlineWithFix(OpBuilder &builder, Operation &op, IRMapping &mapping) { if (auto alloc = dyn_cast(&op)) { auto mapOperand = [&](Value operand, Type expectedType) -> Value { - if (!operand) + if (!operand) { return Value(); + } Value mapped = mapping.lookupOrNull(operand); - if (!mapped) + if (!mapped) { mapped = operand; + } return maybeUnwrapCastToExpected(mapped, expectedType); }; @@ -121,39 +123,48 @@ static void eraseDeadBridgeCasts(func::FuncOp func) { SmallVector deadUnrealized; func.walk([&](UnrealizedConversionCastOp cast) { - if (cast->use_empty()) + if (cast->use_empty()) { deadUnrealized.push_back(cast); + } }); SmallVector deadMemrefCasts; func.walk([&](memref::CastOp cast) { - if (cast->use_empty()) + if (cast->use_empty()) { deadMemrefCasts.push_back(cast); + } }); - if (deadUnrealized.empty() && deadMemrefCasts.empty()) + if (deadUnrealized.empty() && deadMemrefCasts.empty()) { break; + } - for (UnrealizedConversionCastOp cast : llvm::reverse(deadUnrealized)) + for (UnrealizedConversionCastOp cast : llvm::reverse(deadUnrealized)) { cast.erase(); - for (memref::CastOp cast : llvm::reverse(deadMemrefCasts)) + } + for (memref::CastOp cast : llvm::reverse(deadMemrefCasts)) { cast.erase(); + } changed = true; } } static LogicalResult inlineCall(func::CallOp call, func::FuncOp callee) { - if (callee.isExternal()) + if (callee.isExternal()) { return call.emitOpError("callee must have a body before inlining"); + } Block &entry = callee.getBody().front(); - if (entry.getNumArguments() != call.getNumOperands()) + if (entry.getNumArguments() != call.getNumOperands()) { return call.emitOpError("callee argument count mismatch during inlining"); + } auto returnOp = dyn_cast(entry.getTerminator()); - if (!returnOp) + if (!returnOp) { return call.emitOpError("callee must terminate with func.return"); - if (returnOp.getNumOperands() != call.getNumResults()) + } + if (returnOp.getNumOperands() != call.getNumResults()) { return call.emitOpError("callee return/result arity mismatch during inlining"); + } OpBuilder builder(call); IRMapping mapping; @@ -164,10 +175,12 @@ static LogicalResult inlineCall(func::CallOp call, func::FuncOp callee) { for (Operation &op : entry.without_terminator()) { FailureOr handledOr = pto::tryCloneOpLibInlineBridgeOp(builder, op, mapping); - if (failed(handledOr)) + if (failed(handledOr)) { return call.emitOpError("failed to remap OP-Lib inline bridge op"); - if (*handledOr) + } + if (*handledOr) { continue; + } Operation *newOp = cloneOpForInlineWithFix(builder, op, mapping); for (auto [oldRes, newRes] : @@ -178,8 +191,9 @@ static LogicalResult inlineCall(func::CallOp call, func::FuncOp callee) { for (auto [callResult, returnOperand] : llvm::zip(call.getResults(), returnOp.getOperands())) { Value mapped = mapping.lookupOrNull(returnOperand); - if (!mapped) + if (!mapped) { mapped = returnOperand; + } callResult.replaceAllUsesWith(mapped); } @@ -218,25 +232,29 @@ static LogicalResult validateInlineableCalleesHaveBodies( ModuleOp module, InlinePredicate &&shouldInline) { for (ModuleOp funcModule : collectFuncModules(module)) { for (func::FuncOp func : funcModule.getOps()) { - if (func.isExternal() || func.empty()) + if (func.isExternal() || func.empty()) { continue; + } bool failed = false; func.walk([&](func::CallOp call) { auto calleeAttr = call.getCalleeAttr(); - if (!calleeAttr) + if (!calleeAttr) { return; + } func::FuncOp callee = funcModule.lookupSymbol(calleeAttr.getValue()); - if (!callee || !shouldInline(callee) || !callee.isExternal()) + if (!callee || !shouldInline(callee) || !callee.isExternal()) { return; + } emitMissingInstanceBodyError(call, callee); failed = true; }); - if (failed) + if (failed) { return failure(); + } } } @@ -249,12 +267,15 @@ static LogicalResult inlineMatchingCalls( llvm::StringRef debugTag, int &inlinedCalls, int &touchedFuncs) { for (ModuleOp funcModule : collectFuncModules(module)) { for (func::FuncOp func : funcModule.getOps()) { - if (func.isExternal()) + if (func.isExternal()) { continue; - if (isInstanceFunc(func)) + } + if (isInstanceFunc(func)) { continue; - if (func.empty()) + } + if (func.empty()) { continue; + } bool changedThisFunc = false; bool madeProgress = true; @@ -265,17 +286,20 @@ static LogicalResult inlineMatchingCalls( func.walk([&](func::CallOp call) { calls.push_back(call); }); for (func::CallOp oldCall : calls) { - if (!oldCall || !oldCall->getBlock()) + if (!oldCall || !oldCall->getBlock()) { continue; + } auto calleeAttr = oldCall.getCalleeAttr(); - if (!calleeAttr) + if (!calleeAttr) { continue; + } func::FuncOp callee = funcModule.lookupSymbol(calleeAttr.getValue()); - if (!callee || !shouldInline(callee)) + if (!callee || !shouldInline(callee)) { continue; + } if (callee.isExternal()) { oldCall.emitOpError("callee must have a body before inlining"); @@ -304,8 +328,9 @@ static LogicalResult inlineMatchingCalls( oldResult.replaceAllUsesWith(newResult); call.erase(); - if (failed(inlineCall(newCall, callee))) + if (failed(inlineCall(newCall, callee))) { return failure(); + } ++inlinedCalls; changedThisFunc = true; @@ -334,16 +359,20 @@ static void eraseDeadMatchingPrivateFuncs(ModuleOp module, SymbolTable symbolTable(funcModule); SmallVector deadFuncs; for (func::FuncOp func : funcModule.getOps()) { - if (!predicate(func)) + if (!predicate(func)) { continue; - if (func.isPublic()) + } + if (func.isPublic()) { continue; + } auto uses = symbolTable.getSymbolUses(func, funcModule); - if (uses && uses->empty()) + if (uses && uses->empty()) { deadFuncs.push_back(func); + } } - for (func::FuncOp func : deadFuncs) + for (func::FuncOp func : deadFuncs) { func.erase(); + } } } diff --git a/lib/PTO/Transforms/PTOLowerFrontendPipeOpsPass.cpp b/lib/PTO/Transforms/PTOLowerFrontendPipeOpsPass.cpp index 2fb66397e2..1c777b0bf6 100644 --- a/lib/PTO/Transforms/PTOLowerFrontendPipeOpsPass.cpp +++ b/lib/PTO/Transforms/PTOLowerFrontendPipeOpsPass.cpp @@ -64,8 +64,9 @@ static LogicalResult requireFrontendGmSlotBuffer(InitOpT initOp) { template static void propagateFrontendIdAttr(InitOpT initOp, Operation *pipeOp, IRRewriter &rewriter) { - if (!pipeOp) + if (!pipeOp) { return; + } pipeOp->setAttr(kFrontendPipeIdAttrName, rewriter.getI32IntegerAttr(initOp.getId())); } @@ -74,12 +75,14 @@ template static void propagateFixpipePeerKeyAttrs(InitOpT initOp, Operation *pipeOp, IRRewriter &rewriter) { if (!pipeOp || !initOp.getAccPushEpilogueAttr() || - initOp.getDirMask() != kC2VDirMask || !initOp.getC2vConsumerBuf()) + initOp.getDirMask() != kC2VDirMask || !initOp.getC2vConsumerBuf()) { return; + } auto currentFunc = initOp->template getParentOfType(); - if (!currentFunc) + if (!currentFunc) { return; + } auto setPeerKeyAttrs = [&](FlatSymbolRefAttr ownerFuncAttr, StringRef reserveName) { @@ -108,8 +111,9 @@ static void propagateFixpipePeerKeyAttrs(InitOpT initOp, Operation *pipeOp, template static int32_t getFrontendSlotNum(InitOpT initOp) { - if (auto slotNumAttr = initOp.getSlotNumAttr()) + if (auto slotNumAttr = initOp.getSlotNumAttr()) { return slotNumAttr.getInt(); + } return initOp.getDirMask() == kBidirectionalDirMask ? kBidirectionalSlotNum : kSingleDirectionSlotNum; @@ -127,12 +131,14 @@ static std::optional getStaticIndexLikeValue(Value value) { static SmallVector getStaticTensorViewStrides(Value tensor) { SmallVector strides; - if (!tensor) + if (!tensor) { return strides; + } auto makeView = tensor.getDefiningOp(); - if (!makeView) + if (!makeView) { return strides; + } auto tvTy = dyn_cast(makeView.getResult().getType()); if (!tvTy || @@ -142,8 +148,9 @@ static SmallVector getStaticTensorViewStrides(Value tensor) { strides.reserve(makeView.getStrides().size()); for (Value stride : makeView.getStrides()) { auto staticStride = getStaticIndexLikeValue(stride); - if (!staticStride) + if (!staticStride) { return {}; + } strides.push_back(*staticStride); } return strides; @@ -152,8 +159,9 @@ static SmallVector getStaticTensorViewStrides(Value tensor) { static void propagateGlobalTensorStrides(DeclareGlobalOp decl, ArrayRef strides, IRRewriter &rewriter) { - if (strides.empty()) + if (strides.empty()) { return; + } decl->setAttr(kGlobalTensorStridesAttrName, rewriter.getDenseI64ArrayAttr(strides)); } @@ -176,8 +184,9 @@ static FailureOr createFrontendGlobalTensorPipe(InitOpT initOp, IntegerAttr localSlotNumAttr; if (localAddr) { localSlotNumAttr = initOp.getLocalSlotNumAttr(); - if (!localSlotNumAttr) + if (!localSlotNumAttr) { localSlotNumAttr = rewriter.getI32IntegerAttr(slotNum); + } } auto pipe = rewriter.create( loc, pipeTy, dirAttr, slotSizeAttr, slotNumAttr, localSlotNumAttr, @@ -217,13 +226,15 @@ static FailureOr createFrontendLocalPipe(InitOpT initOp, if (failed(requireFrontendGmSlotBuffer(initOp))) return failure(); - if (!localAddr) + if (!localAddr) { return initOp.emitOpError( "requires local consumer buffer operands for local FIFO pipe lowering"); + } IntegerAttr localSlotNumAttr = initOp.getLocalSlotNumAttr(); - if (!localSlotNumAttr) + if (!localSlotNumAttr) { localSlotNumAttr = rewriter.getI32IntegerAttr(slotNum); + } auto pipe = rewriter.create( loc, pipeTy, dirAttr, slotSizeAttr, slotNumAttr, localSlotNumAttr, IntegerAttr{}, noSplitAttr, accPushEpilogueAttr, initOp.getGmSlotBuffer(), @@ -317,21 +328,25 @@ template static void propagateFrontendNoSplitAttr(InitOpT initOp, const FrontendPipeHandles &handles) { auto noSplitAttr = initOp.getNosplitAttr(); - if (!noSplitAttr) + if (!noSplitAttr) { return; + } - if (handles.anchorOp) + if (handles.anchorOp) { handles.anchorOp->setAttr("nosplit", noSplitAttr); + } Operation *c2vOp = handles.c2vPipe ? handles.c2vPipe.getDefiningOp() : nullptr; Operation *v2cOp = handles.v2cPipe ? handles.v2cPipe.getDefiningOp() : nullptr; - if (c2vOp && c2vOp != handles.anchorOp) + if (c2vOp && c2vOp != handles.anchorOp) { c2vOp->setAttr("nosplit", noSplitAttr); - if (v2cOp && v2cOp != handles.anchorOp && v2cOp != c2vOp) + } + if (v2cOp && v2cOp != handles.anchorOp && v2cOp != c2vOp) { v2cOp->setAttr("nosplit", noSplitAttr); + } } template @@ -383,8 +398,9 @@ static FailureOr lowerInitIfPresent(func::FuncOp funcOp, return WalkResult::advance(); }); - if (hasDuplicateId) + if (hasDuplicateId) { return failure(); + } if (hasAicInit && hasAivInit) { funcOp.emitOpError("cannot mix pto.aic_initialize_pipe and " diff --git a/lib/PTO/Transforms/PTOLowerToOpLibCalls.cpp b/lib/PTO/Transforms/PTOLowerToOpLibCalls.cpp index dde973d6c7..0c8d19d7fd 100644 --- a/lib/PTO/Transforms/PTOLowerToOpLibCalls.cpp +++ b/lib/PTO/Transforms/PTOLowerToOpLibCalls.cpp @@ -19,12 +19,14 @@ FailureOr mlir::pto::tryCloneOpLibInlineBridgeOp(OpBuilder &builder, Operation &op, IRMapping &mapping) { if (auto cast = dyn_cast(&op)) { - if (cast->getNumOperands() != 1 || cast->getNumResults() != 1) + if (cast->getNumOperands() != 1 || cast->getNumResults() != 1) { return failure(); + } Value mappedSrc = mapping.lookupOrNull(cast.getOperand(0)); - if (!mappedSrc) + if (!mappedSrc) { return failure(); + } Type dstTy = cast.getResult(0).getType(); if (mappedSrc.getType() == dstTy) { diff --git a/lib/PTO/Transforms/PTOMaterializeSIMTPersistentFragment.cpp b/lib/PTO/Transforms/PTOMaterializeSIMTPersistentFragment.cpp index 6aba8b1640..c7d4490d48 100644 --- a/lib/PTO/Transforms/PTOMaterializeSIMTPersistentFragment.cpp +++ b/lib/PTO/Transforms/PTOMaterializeSIMTPersistentFragment.cpp @@ -309,8 +309,9 @@ buildPersistentTransformWorklist(const PersistentMaterializationPlan &plan, accessLane.op ? accessLane.op->getParentOfType() : pto::SectionSimtOp(); - if (accessSection == section) - localAccesses.push_back(accessLane); +if (accessSection == section) { + localAccesses.push_back(accessLane); + } } sectionWorklist.elements.push_back( {&fragment, &residentElement, std::move(localAccesses)}); @@ -412,8 +413,9 @@ static LogicalResult materializeSection(const PersistentSectionWorklist §ionWorklist, const DataLayout &dataLayout, DominanceInfo &dominance) { const auto &elements = sectionWorklist.elements; - if (elements.empty()) + if (elements.empty()) { return success(); + } pto::SectionSimtOp section = sectionWorklist.section; Block &body = section.getBody().front(); @@ -439,8 +441,9 @@ materializeSection(const PersistentSectionWorklist §ionWorklist, const PersistentFragmentAnalysis &fragment = *element.fragment; const ResidentElementPlan &residentElement = *element.residentElement; LLVM::AllocaOp allocaOp = fragment.allocaOp; - if (section == fragment.initSection) + if (section == fragment.initSection) { continue; + } rewrites[elementIndex].resumeValue = entryBuilder @@ -453,8 +456,9 @@ materializeSection(const PersistentSectionWorklist §ionWorklist, // its accesses. Carry sections seed it from resume; init sections rely on // the previously validated first-store initialization. for (auto [elementIndex, element] : llvm::enumerate(elements)) { - if (element.accesses.empty()) + if (element.accesses.empty()) { continue; + } const PersistentFragmentAnalysis &fragment = *element.fragment; LLVM::AllocaOp allocaOp = fragment.allocaOp; @@ -477,8 +481,9 @@ materializeSection(const PersistentSectionWorklist §ionWorklist, size_t rewrittenAccessCount = 0; for (Operation &op : llvm::make_early_inc_range(body)) { auto accessIt = laneRewritesByAccess.find(&op); - if (accessIt == laneRewritesByAccess.end()) + if (accessIt == laneRewritesByAccess.end()) { continue; + } if (failed(rewritePersistentAccess(&op, accessIt->second, rewrites))) return failure(); ++rewrittenAccessCount; @@ -523,8 +528,9 @@ materializeSection(const PersistentSectionWorklist §ionWorklist, promotionBuilder.setInsertionPointToStart(&body); for (auto [elementIndex, element] : llvm::enumerate(elements)) { LLVM::AllocaOp proxy = rewrites[elementIndex].proxy; - if (!proxy) + if (!proxy) { continue; + } SmallVector allocators{ cast(proxy.getOperation())}; if (failed(tryToPromoteMemorySlots(allocators, promotionBuilder, dataLayout, @@ -536,8 +542,9 @@ materializeSection(const PersistentSectionWorklist §ionWorklist, } } - if (proxyArraySize && proxyArraySize.use_empty()) + if (proxyArraySize && proxyArraySize.use_empty()) { proxyArraySize.getDefiningOp()->erase(); + } return success(); } diff --git a/lib/PTO/Transforms/PTOMaterializeTileOpSections.cpp b/lib/PTO/Transforms/PTOMaterializeTileOpSections.cpp index 6c2a17c11a..efaacfdf28 100644 --- a/lib/PTO/Transforms/PTOMaterializeTileOpSections.cpp +++ b/lib/PTO/Transforms/PTOMaterializeTileOpSections.cpp @@ -120,8 +120,9 @@ static std::optional traceToFunctionArgument(Value value, static void applyEffectToAllTileArguments(func::FuncOp function, uint8_t effect, SmallVectorImpl &effects) { for (auto [index, type] : llvm::enumerate(function.getArgumentTypes())) - if (isa(type)) + if (isa(type)) { effects[index] |= effect; + } } static SmallVector @@ -141,8 +142,9 @@ collectDirectArgumentEffects(func::FuncOp function) { effect = ReadEffect; else if (isa(instance.getEffect())) effect = WriteEffect; - if (effect == NoEffect || !instance.getValue()) + if (effect == NoEffect || !instance.getValue()) { continue; + } if (auto argument = traceToFunctionArgument(instance.getValue(), function)) @@ -168,8 +170,9 @@ static void summarizeSimtLaunchEffects(func::FuncOp helper, SimtLaunchOp launch, SmallVector calleeEffects = collectDirectArgumentEffects(callee); for (auto [argument, calleeEffect] : llvm::zip_equal(launch.getArgs(), calleeEffects)) { - if (calleeEffect == NoEffect) + if (calleeEffect == NoEffect) { continue; + } if (auto helperArgument = traceToFunctionArgument(argument, helper)) effects[*helperArgument] |= calleeEffect; else @@ -206,8 +209,9 @@ static bool addValidShapeRequirement(ValidShapeRequirements &requirements, func::FuncOp function, unsigned argumentIndex) { auto &indices = requirements[function.getOperation()]; - if (llvm::is_contained(indices, argumentIndex)) + if (llvm::is_contained(indices, argumentIndex)) { return false; + } indices.push_back(argumentIndex); llvm::sort(indices); return true; @@ -238,8 +242,9 @@ collectTileOpValidShapeRequirements(func::FuncOp helper, } unsigned dimension = isa(op) ? 0 : 1; - if (tileType.getValidShape()[dimension] < 0) + if (tileType.getValidShape()[dimension] < 0) { addValidShapeRequirement(requirements, helper, argument.getArgNumber()); + } return WalkResult::advance(); }); return status; @@ -257,11 +262,13 @@ propagateValidShapeRequirements(ModuleOp module, for (func::CallOp call : calls) { auto callee = SymbolTable::lookupNearestSymbolFrom( call.getOperation(), call.getCalleeAttr()); - if (!callee) + if (!callee) { continue; + } auto required = requirements.find(callee.getOperation()); - if (required == requirements.end()) + if (required == requirements.end()) { continue; + } SmallVector requiredArguments(required->second); auto caller = call->getParentOfType(); @@ -270,17 +277,20 @@ propagateValidShapeRequirements(ModuleOp module, "cannot propagate Tile valid-shape metadata without a caller " "function"); for (unsigned calleeArgument : requiredArguments) { - if (calleeArgument >= call.getNumOperands()) + if (calleeArgument >= call.getNumOperands()) { return call.emitOpError( "TileOp call has fewer operands than its helper ABI"); + } auto callerArgument = traceToFunctionArgument(call.getOperand(calleeArgument), caller); - if (!callerArgument) + if (!callerArgument) { continue; + } auto tileType = dyn_cast( caller.getArgument(*callerArgument).getType()); - if (!tileType || !tileType.hasDynamicValid()) + if (!tileType || !tileType.hasDynamicValid()) { continue; + } changed |= addValidShapeRequirement(requirements, caller, *callerArgument); } @@ -335,8 +345,9 @@ static std::optional> resolveCallValidShape(Value tile, Operation *anchor, func::FuncOp caller, const ExpandedValidShapeArguments &expandedArguments, OpBuilder &builder) { - if (!tile) + if (!tile) { return std::nullopt; + } auto tileType = dyn_cast(tile.getType()); if (tileType && tileType.getValidShape().size() == 2 && @@ -357,8 +368,9 @@ resolveCallValidShape(Value tile, Operation *anchor, func::FuncOp caller, } } - if (!tileType || tileType.getValidShape().size() != 2) + if (!tileType || tileType.getValidShape().size() != 2) { return std::nullopt; + } // set_validshape mutates the Tile metadata in place. Reading it at the call // preserves the executed update across sequential and structured control @@ -378,19 +390,22 @@ static LogicalResult expandValidShapeCallOperands( for (func::CallOp call : calls) { auto callee = SymbolTable::lookupNearestSymbolFrom( call.getOperation(), call.getCalleeAttr()); - if (!callee) + if (!callee) { continue; + } auto required = requirements.find(callee.getOperation()); - if (required == requirements.end()) + if (required == requirements.end()) { continue; + } auto caller = call->getParentOfType(); OpBuilder builder(call); SmallVector metadataOperands; for (unsigned argumentIndex : required->second) { - if (argumentIndex >= call.getNumOperands()) + if (argumentIndex >= call.getNumOperands()) { return call.emitOpError( "TileOp call has fewer operands than its helper ABI"); + } auto metadata = resolveCallValidShape(call.getOperand(argumentIndex), call, caller, expandedArguments, builder); @@ -479,14 +494,16 @@ materializeTileOpValidShapeABI(ModuleOp module, static LogicalResult verifyTileOpABI(func::FuncOp helper) { for (auto [index, type] : llvm::enumerate(helper.getArgumentTypes())) { - if (!isTileOrScalarType(type)) + if (!isTileOrScalarType(type)) { return helper.emitOpError() << "tileop argument #" << index << " must be !pto.tile_buf or a PTO scalar, got " << type; + } } - if (helper.getNumResults() != 0) + if (helper.getNumResults() != 0) { return helper.emitOpError("tileop helpers must not return values; write " "results through mutable Tile parameters"); + } return success(); } @@ -584,8 +601,9 @@ static LogicalResult inferTileOpKind(func::FuncOp helper, return WalkResult::advance(); }); - if (failed(status)) + if (failed(status)) { return failure(); + } if (firstVector && firstCube) { InFlightDiagnostic diag = helper.emitOpError( "mixes Vector and Cube compute operations in one tileop helper"); @@ -593,9 +611,10 @@ static LogicalResult inferTileOpKind(func::FuncOp helper, diag.attachNote(firstCube->getLoc()) << "first Cube operation is here"; return failure(); } - if (!firstVector && !firstCube) + if (!firstVector && !firstCube) { return helper.emitOpError("contains no Vector or Cube compute operation " "from which to infer tileop kind"); + } kind = firstVector ? PhysicalSectionKind::Vector : PhysicalSectionKind::Cube; return success(); } @@ -610,10 +629,12 @@ static LogicalResult materializeTileOpSection(func::FuncOp helper, return helper.emitOpError("requires a func.return terminator"); SmallVector roots; - for (Operation &op : entry.without_terminator()) + for (Operation &op : entry.without_terminator()) { roots.push_back(&op); - if (roots.empty()) + } + if (roots.empty()) { return helper.emitOpError("contains no materializable compute body"); + } OpBuilder builder(roots.front()); Operation *sectionOperation = @@ -626,8 +647,9 @@ static LogicalResult materializeTileOpSection(func::FuncOp helper, auto *sectionBlock = new Block(); sectionBody.push_back(sectionBlock); - for (Operation *root : roots) + for (Operation *root : roots) { root->moveBefore(sectionBlock, sectionBlock->end()); + } helper->setAttr( kTileOpKindAttr, StringAttr::get(helper.getContext(), diff --git a/lib/PTO/Transforms/PTONarrowVPTOLoopCounters.cpp b/lib/PTO/Transforms/PTONarrowVPTOLoopCounters.cpp index da25a77941..561397a9d5 100644 --- a/lib/PTO/Transforms/PTONarrowVPTOLoopCounters.cpp +++ b/lib/PTO/Transforms/PTONarrowVPTOLoopCounters.cpp @@ -118,8 +118,9 @@ struct NarrowVecScopeLoopCounterPattern : public OpRewritePattern { Block *oldBody = forOp.getBody(); Block *newBody = newFor.getBody(); - if (!newBody->empty()) + if (!newBody->empty()) { rewriter.eraseOp(newBody->getTerminator()); + } rewriter.setInsertionPointToStart(newBody); Value restoredInductionVar = restoreInductionVariableType( diff --git a/lib/PTO/Transforms/PTONormalizeUncoveredTileSections.cpp b/lib/PTO/Transforms/PTONormalizeUncoveredTileSections.cpp index 057d38a07e..7879737662 100644 --- a/lib/PTO/Transforms/PTONormalizeUncoveredTileSections.cpp +++ b/lib/PTO/Transforms/PTONormalizeUncoveredTileSections.cpp @@ -43,13 +43,16 @@ struct UncoveredTopLevelSegment { static void mergeSegmentSummary(UncoveredTopLevelSegment &dst, const UncoveredTopLevelSegment &src) { - if (!src.firstOp) + if (!src.firstOp) { return; - if (!dst.firstOp) + } + if (!dst.firstOp) { dst.firstOp = src.firstOp; + } dst.lastOp = src.lastOp; - if (!dst.firstTileCarrierOp) + if (!dst.firstTileCarrierOp) { dst.firstTileCarrierOp = src.firstTileCarrierOp; + } dst.containsTileOp |= src.containsTileOp; dst.containsNestedExplicitSection |= src.containsNestedExplicitSection; dst.vectorTileOpCount += src.vectorTileOpCount; @@ -63,8 +66,9 @@ static bool isExplicitSection(Operation *op) { } static bool isTileLikeOp(Operation *op) { - if (!op) + if (!op) { return false; + } return isa(op) && op->getName().getStringRef().starts_with("pto.t"); } @@ -73,8 +77,9 @@ static bool isTileLikeOp(Operation *op) { // section inference conservative: only operations with an unambiguous engine // ownership are treated as section carriers. static bool isRawSectionCarrierOp(Operation *op) { - if (isa(op) || isa(op)) + if (isa(op) || isa(op)) { return false; + } return op && isa(op); @@ -112,31 +117,39 @@ static std::optional classifyMteOpByAddressSpace(MteOpInterface mteOp) { Value sourceValue = mteOp.getSource(); std::optional source; - if (sourceValue) + if (sourceValue) { source = getBufferAddressSpace(sourceValue.getType()); + } Value destinationValue = mteOp.getDestination(); if (auto ptoDpsOp = dyn_cast(mteOp.getOperation())) { OperandRange inits = ptoDpsOp.getDpsInits(); - if (!inits.empty()) + if (!inits.empty()) { destinationValue = inits.front(); + } } std::optional destination; - if (destinationValue) + if (destinationValue) { destination = getBufferAddressSpace(destinationValue.getType()); + } - if (source && *source == AddressSpace::ACC) + if (source && *source == AddressSpace::ACC) { return InferredSectionKind::Cube; - if (source && *source == AddressSpace::VEC) + } + if (source && *source == AddressSpace::VEC) { return InferredSectionKind::Vector; - if (!destination) + } + if (!destination) { return std::nullopt; + } if (*destination == AddressSpace::MAT || *destination == AddressSpace::LEFT || *destination == AddressSpace::RIGHT || *destination == AddressSpace::BIAS || - *destination == AddressSpace::SCALING) + *destination == AddressSpace::SCALING) { return InferredSectionKind::Cube; - if (*destination == AddressSpace::VEC) + } + if (*destination == AddressSpace::VEC) { return InferredSectionKind::Vector; + } return std::nullopt; } @@ -144,10 +157,12 @@ static std::optional classifyRawSectionCarrierOp(Operation *op) { if (!isRawSectionCarrierOp(op)) return std::nullopt; - if (isa(op)) + if (isa(op)) { return InferredSectionKind::Vector; - if (isa(op)) + } + if (isa(op)) { return InferredSectionKind::Cube; + } if (auto mteOp = dyn_cast(op)) { if (auto kind = classifyMteOpByAddressSpace(mteOp)) return kind; @@ -158,16 +173,18 @@ classifyRawSectionCarrierOp(Operation *op) { // therefore determines physical ownership even when the consumer pipe is // shared; a shared source remains ambiguous and must not inherit ownership // from its peer. - if (isSharedSyncPipe(setFlag.getSrcPipe().getPipe())) + if (isSharedSyncPipe(setFlag.getSrcPipe().getPipe())) { return std::nullopt; + } return classifySyncPipe(setFlag.getSrcPipe().getPipe()); } if (auto waitFlag = dyn_cast(op)) { // wait_flag executes on the consumer (destination) pipe. Mirror the // producer rule above instead of rejecting a uniquely owned consumer just // because its producer uses a shared pipe. - if (isSharedSyncPipe(waitFlag.getDstPipe().getPipe())) + if (isSharedSyncPipe(waitFlag.getDstPipe().getPipe())) { return std::nullopt; + } return classifySyncPipe(waitFlag.getDstPipe().getPipe()); } if (auto syncSet = dyn_cast(op)) @@ -186,15 +203,18 @@ static bool isRawVPTOVectorTransientType(Type type) { } static bool isRawVPTOVectorLikeOp(Operation *op) { - if (!op) + if (!op) { return false; + } for (Value operand : op->getOperands()) { - if (isRawVPTOVectorTransientType(operand.getType())) + if (isRawVPTOVectorTransientType(operand.getType())) { return true; + } } for (Value result : op->getResults()) { - if (isRawVPTOVectorTransientType(result.getType())) + if (isRawVPTOVectorTransientType(result.getType())) { return true; + } } return false; } @@ -248,13 +268,15 @@ static std::optional getBufferAddressSpace(Type type) { static void collectTileAddressSpaces(Type type, SmallVectorImpl &spaces) { - if (std::optional addressSpace = getBufferAddressSpace(type)) + if (std::optional addressSpace = getBufferAddressSpace(type)) { spaces.push_back(*addressSpace); + } } static std::optional getPipeHandleDirMask(Value pipeHandle) { - if (!pipeHandle) + if (!pipeHandle) { return std::nullopt; + } if (auto init = pipeHandle.getDefiningOp()) return init.getDirMask(); if (auto init = pipeHandle.getDefiningOp()) @@ -286,50 +308,63 @@ static std::optional classifyInternalPipeTileOp(Operation *op) { if (auto push = dyn_cast(op)) { std::optional dirMask = getPipeHandleDirMask(push.getPipeHandle()); - if (!dirMask) + if (!dirMask) { return std::nullopt; - if (*dirMask == 1) + } + if (*dirMask == 1) { return InferredSectionKind::Cube; - if (*dirMask == 2) + } + if (*dirMask == 2) { return InferredSectionKind::Vector; + } return classifyTileSectionByAddressSpace( getBufferAddressSpace(push.getTile().getType())); } if (auto pop = dyn_cast(op)) { std::optional dirMask = getPipeHandleDirMask(pop.getPipeHandle()); - if (!dirMask) + if (!dirMask) { return std::nullopt; - if (*dirMask == 1) + } + if (*dirMask == 1) { return InferredSectionKind::Vector; - if (*dirMask == 2) + } + if (*dirMask == 2) { return InferredSectionKind::Cube; + } return classifyTileSectionByAddressSpace( getBufferAddressSpace(pop.getTile().getType())); } if (auto free = dyn_cast(op)) { std::optional dirMask = getPipeHandleDirMask(free.getPipeHandle()); - if (!dirMask) + if (!dirMask) { return std::nullopt; - if (*dirMask == 1) + } + if (*dirMask == 1) { return InferredSectionKind::Vector; - if (*dirMask == 2) + } + if (*dirMask == 2) { return InferredSectionKind::Cube; - if (!free.getEntry()) + } + if (!free.getEntry()) { return std::nullopt; + } return classifyTileSectionByAddressSpace( getBufferAddressSpace(free.getEntry().getType())); } if (auto alloc = dyn_cast(op)) { std::optional dirMask = getPipeHandleDirMask(alloc.getPipeHandle()); - if (!dirMask) + if (!dirMask) { return std::nullopt; - if (*dirMask == 1) + } + if (*dirMask == 1) { return InferredSectionKind::Cube; - if (*dirMask == 2) + } + if (*dirMask == 2) { return InferredSectionKind::Vector; + } return std::nullopt; } @@ -386,25 +421,30 @@ classifyTileOpByAddressSpace(Operation *op) { } } - if (sawCubeOnly) + if (sawCubeOnly) { return InferredSectionKind::Cube; - if (sawVec) + } + if (sawVec) { return InferredSectionKind::Vector; - if (sawMat) + } + if (sawMat) { return classifyTileOpByPipe(op); + } return std::nullopt; } static std::optional classifyTLoadByDestinationAddressSpace(Operation *op) { - if (!isa(op)) + if (!isa(op)) { return std::nullopt; + } auto tload = cast(op); std::optional dstSpace = getBufferAddressSpace(tload.getDst().getType()); - if (!dstSpace) + if (!dstSpace) { return std::nullopt; + } switch (*dstSpace) { case AddressSpace::VEC: @@ -423,14 +463,16 @@ classifyTLoadByDestinationAddressSpace(Operation *op) { static std::optional classifyTStoreBySourceAddressSpace(Operation *op) { - if (!isa(op)) + if (!isa(op)) { return std::nullopt; + } auto tstore = cast(op); std::optional srcSpace = getBufferAddressSpace(tstore.getSrc().getType()); - if (!srcSpace) + if (!srcSpace) { return std::nullopt; + } switch (*srcSpace) { case AddressSpace::VEC: @@ -448,19 +490,24 @@ classifyTStoreBySourceAddressSpace(Operation *op) { } static std::optional classifyTileOp(Operation *op) { - if (std::optional kind = classifyTileOpByName(op)) + if (std::optional kind = classifyTileOpByName(op)) { return kind; - if (std::optional kind = classifyInternalPipeTileOp(op)) + } + if (std::optional kind = classifyInternalPipeTileOp(op)) { return kind; + } if (std::optional kind = - classifyTLoadByDestinationAddressSpace(op)) + classifyTLoadByDestinationAddressSpace(op)) { return kind; + } if (std::optional kind = - classifyTStoreBySourceAddressSpace(op)) + classifyTStoreBySourceAddressSpace(op)) { return kind; + } if (std::optional kind = - classifyTileOpByAddressSpace(op)) + classifyTileOpByAddressSpace(op)) { return kind; + } return classifyTileOpByPipe(op); } @@ -479,12 +526,15 @@ enum class FunctionKindCacheState : uint8_t { static void inspectModuleKindOperation(Operation *op, ModuleKindSummary &summary) { - if (!op) + if (!op) { return; - if (isa(op) || isa(op)) + } + if (isa(op) || isa(op)) { ++summary.vectorCount; - if (isExplicitSection(op)) + } + if (isExplicitSection(op)) { return; + } if (isRawSectionCarrierOp(op)) { if (std::optional kind = @@ -540,8 +590,9 @@ decodeFunctionKind(FunctionKindCacheState state) { } static func::CallOp getTransparentWrapperCall(func::FuncOp funcOp) { - if (!funcOp || funcOp.isDeclaration() || !funcOp.getBody().hasOneBlock()) + if (!funcOp || funcOp.isDeclaration() || !funcOp.getBody().hasOneBlock()) { return nullptr; + } Block &entryBlock = funcOp.getBody().front(); func::CallOp callOp; @@ -551,21 +602,26 @@ static func::CallOp getTransparentWrapperCall(func::FuncOp funcOp) { returnOp = ret; continue; } - if (callOp) + if (callOp) { return nullptr; + } callOp = dyn_cast(op); - if (!callOp) + if (!callOp) { return nullptr; + } } - if (!callOp || !returnOp) + if (!callOp || !returnOp) { return nullptr; - if (returnOp.getNumOperands() != callOp.getNumResults()) + } + if (returnOp.getNumOperands() != callOp.getNumResults()) { return nullptr; + } for (auto [returned, forwarded] : llvm::zip(returnOp.getOperands(), callOp.getResults())) { - if (returned != forwarded) + if (returned != forwarded) { return nullptr; + } } return callOp; } @@ -573,13 +629,15 @@ static func::CallOp getTransparentWrapperCall(func::FuncOp funcOp) { static std::optional inferWholeFunctionKind( func::FuncOp funcOp, llvm::DenseMap &cache) { - if (!funcOp || funcOp.isDeclaration()) + if (!funcOp || funcOp.isDeclaration()) { return std::nullopt; + } auto cacheIt = cache.find(funcOp.getOperation()); if (cacheIt != cache.end()) { - if (cacheIt->second == FunctionKindCacheState::InProgress) + if (cacheIt->second == FunctionKindCacheState::InProgress) { return std::nullopt; + } return decodeFunctionKind(cacheIt->second); } cache[funcOp.getOperation()] = FunctionKindCacheState::InProgress; @@ -589,10 +647,11 @@ static std::optional inferWholeFunctionKind( std::optional inferredKind; if (summary.ambiguousOps.empty() && !(summary.vectorCount && summary.cubeCount)) { - if (summary.vectorCount) + if (summary.vectorCount) { inferredKind = InferredSectionKind::Vector; - else if (summary.cubeCount) + } else if (summary.cubeCount) { inferredKind = InferredSectionKind::Cube; + } } if (!inferredKind) { @@ -618,8 +677,9 @@ static void assignModuleKernelKind(ModuleOp module, InferredSectionKind kind) { static void assignFunctionKernelKind(func::FuncOp funcOp, InferredSectionKind kind) { - if (!funcOp) + if (!funcOp) { return; + } FunctionKernelKind kernelKind = kind == InferredSectionKind::Vector ? FunctionKernelKind::Vector @@ -629,50 +689,59 @@ static void assignFunctionKernelKind(func::FuncOp funcOp, } static LogicalResult tryAssignWholeModuleKernelKind(ModuleOp module) { - if (!module || module->hasAttr(FunctionKernelKindAttr::name)) + if (!module || module->hasAttr(FunctionKernelKindAttr::name)) { return success(); + } SmallVector defs; for (auto funcOp : module.getOps()) { - if (!funcOp.isDeclaration()) + if (!funcOp.isDeclaration()) { defs.push_back(funcOp); + } } - if (defs.empty()) + if (defs.empty()) { return success(); + } llvm::DenseMap cache; std::optional commonKind; for (func::FuncOp funcOp : defs) { - if (hasAnySection(funcOp)) + if (hasAnySection(funcOp)) { return success(); + } std::optional funcKind = inferWholeFunctionKind(funcOp, cache); - if (!funcKind) + if (!funcKind) { return success(); + } if (!commonKind) { commonKind = funcKind; continue; } - if (*commonKind != *funcKind) + if (*commonKind != *funcKind) { return success(); + } } - if (!commonKind) + if (!commonKind) { return success(); + } assignModuleKernelKind(module, *commonKind); return success(); } static LogicalResult tryAssignWholeFunctionKernelKind(func::FuncOp funcOp) { if (!funcOp || funcOp.isDeclaration() || hasAnySection(funcOp) || - hasKnownKernelKindContext(funcOp)) + hasKnownKernelKindContext(funcOp)) { return success(); + } llvm::DenseMap cache; std::optional kind = inferWholeFunctionKind(funcOp, cache); - if (!kind) + if (!kind) { return success(); + } assignFunctionKernelKind(funcOp, *kind); return success(); @@ -680,8 +749,9 @@ static LogicalResult tryAssignWholeFunctionKernelKind(func::FuncOp funcOp) { static void inspectSegmentOperation(Operation *op, UncoveredTopLevelSegment &segment) { - if (!op) + if (!op) { return; + } if (isTileLikeOp(op) || isRawSectionCarrierOp(op)) { segment.containsTileOp = true; @@ -712,14 +782,18 @@ static void inspectSegmentOperation(Operation *op, static std::optional inferSegmentKind(const UncoveredTopLevelSegment &segment) { - if (!segment.ambiguousTileOps.empty()) + if (!segment.ambiguousTileOps.empty()) { return std::nullopt; - if (segment.vectorTileOpCount && segment.cubeTileOpCount) + } + if (segment.vectorTileOpCount && segment.cubeTileOpCount) { return std::nullopt; - if (segment.vectorTileOpCount) + } + if (segment.vectorTileOpCount) { return InferredSectionKind::Vector; - if (segment.cubeTileOpCount) + } + if (segment.cubeTileOpCount) { return InferredSectionKind::Cube; + } return std::nullopt; } @@ -731,8 +805,9 @@ static UncoveredTopLevelSegment summarizeTopLevelOperation(Operation *op) { summary.firstOp = op; summary.lastOp = op; inspectSegmentOperation(op, summary); - if (summary.containsTileOp) + if (summary.containsTileOp) { summary.firstTileCarrierOp = op; + } return summary; } @@ -745,8 +820,9 @@ static void collectUncoveredTopLevelSegments( UncoveredTopLevelSegment current; auto flushCurrent = [&]() { - if (!current.firstOp) + if (!current.firstOp) { return; + } segments.push_back(current); current = {}; }; @@ -797,8 +873,9 @@ wrapUncoveredTopLevelSegment(func::FuncOp funcOp, Block &entryBlock = funcOp.getBody().front(); Operation *firstOp = segment.firstOp; Operation *lastOp = segment.lastOp; - if (!firstOp || !lastOp) + if (!firstOp || !lastOp) { return; + } OpBuilder builder(firstOp); auto sectionOp = builder.create(firstOp->getLoc()); @@ -854,18 +931,21 @@ emitResidualUncoveredTileSegmentError(func::FuncOp funcOp, } static LogicalResult normalizeFunction(func::FuncOp funcOp) { - if (hasKnownKernelKindContext(funcOp)) + if (hasKnownKernelKindContext(funcOp)) { return success(); + } SmallVector segments; collectUncoveredTopLevelSegments(funcOp, segments); for (const UncoveredTopLevelSegment &segment : llvm::reverse(segments)) { - if (!segment.containsTileOp || segment.containsNestedExplicitSection) + if (!segment.containsTileOp || segment.containsNestedExplicitSection) { continue; + } std::optional kind = inferSegmentKind(segment); - if (!kind) + if (!kind) { return emitSegmentInferenceError(funcOp, segment); + } switch (*kind) { case InferredSectionKind::Cube: @@ -881,14 +961,16 @@ static LogicalResult normalizeFunction(func::FuncOp funcOp) { static LogicalResult verifyFunctionHasNoResidualUncoveredTileSegments(func::FuncOp funcOp) { - if (hasKnownKernelKindContext(funcOp)) + if (hasKnownKernelKindContext(funcOp)) { return success(); + } SmallVector segments; collectUncoveredTopLevelSegments(funcOp, segments); for (const UncoveredTopLevelSegment &segment : segments) { - if (!segment.containsTileOp) + if (!segment.containsTileOp) { continue; + } return emitResidualUncoveredTileSegmentError(funcOp, segment); } return success(); diff --git a/lib/PTO/Transforms/PTOOutlineSIMTSections.cpp b/lib/PTO/Transforms/PTOOutlineSIMTSections.cpp index 90df3a527f..c34174ba2d 100644 --- a/lib/PTO/Transforms/PTOOutlineSIMTSections.cpp +++ b/lib/PTO/Transforms/PTOOutlineSIMTSections.cpp @@ -43,12 +43,14 @@ using namespace mlir; namespace { static bool isDefinedInside(Operation *scope, Value value) { - if (Operation *defOp = value.getDefiningOp()) + if (Operation *defOp = value.getDefiningOp()) { return scope->isAncestor(defOp); + } auto blockArg = dyn_cast(value); - if (!blockArg) + if (!blockArg) { return false; + } Operation *owner = blockArg.getOwner()->getParentOp(); return owner && scope->isAncestor(owner); @@ -61,14 +63,17 @@ static LogicalResult collectCaptures(pto::SectionSimtOp sectionOp, sectionOp.getBody().walk([&](Operation *op) { for (Value operand : op->getOperands()) { - if (isDefinedInside(scope, operand)) + if (isDefinedInside(scope, operand)) { continue; + } if (Operation *defOp = operand.getDefiningOp()) { - if (defOp->hasTrait()) + if (defOp->hasTrait()) { continue; + } } - if (seen.insert(operand).second) + if (seen.insert(operand).second) { captures.push_back(operand); + } } }); @@ -85,21 +90,25 @@ static void cloneExternalConstants(pto::SectionSimtOp sectionOp, for (Value operand : op->getOperands()) { Operation *defOp = operand.getDefiningOp(); if (!defOp || isDefinedInside(scope, operand) || - !defOp->hasTrait()) + !defOp->hasTrait()) { continue; - if (seen.insert(defOp).second) + } + if (seen.insert(defOp).second) { constants.push_back(defOp); + } } }); - for (Operation *constant : constants) + for (Operation *constant : constants) { builder.clone(*constant, mapping); + } } static LogicalResult verifySectionCanBeOutlined(pto::SectionSimtOp sectionOp) { func::FuncOp parentFunc = sectionOp->getParentOfType(); - if (!parentFunc) + if (!parentFunc) { return sectionOp.emitOpError("must be nested in a func.func"); + } if (parentFunc->hasAttr(pto::kPTOSimtEntryAttrName)) { return sectionOp.emitOpError() @@ -107,11 +116,13 @@ static LogicalResult verifySectionCanBeOutlined(pto::SectionSimtOp sectionOp) { << pto::kPTOSimtEntryAttrName << "'"; } - if (!sectionOp.getBody().hasOneBlock()) + if (!sectionOp.getBody().hasOneBlock()) { return sectionOp.emitOpError("requires a single-block body"); + } - if (sectionOp.getBody().front().getNumArguments() != 0) + if (sectionOp.getBody().front().getNumArguments() != 0) { return sectionOp.emitOpError("does not support region block arguments"); + } bool hasNestedSection = false; sectionOp.getBody().walk([&](pto::SectionSimtOp nested) { @@ -121,15 +132,17 @@ static LogicalResult verifySectionCanBeOutlined(pto::SectionSimtOp sectionOp) { } return WalkResult::advance(); }); - if (hasNestedSection) + if (hasNestedSection) { return sectionOp.emitOpError("does not support nested pto.section.simt"); + } Operation *scope = sectionOp.getOperation(); WalkResult escapeCheck = sectionOp.getBody().walk([&](Operation *op) { for (Value result : op->getResults()) { for (Operation *user : result.getUsers()) { - if (!scope->isAncestor(user)) + if (!scope->isAncestor(user)) { return WalkResult::interrupt(); + } } } return WalkResult::advance(); @@ -149,8 +162,9 @@ static std::string getUniqueHelperName(ModuleOp module, func::FuncOp parentFunc, do { std::string candidate = (Twine(parentName) + "_simt_" + Twine(outlineIndex++)).str(); - if (!module.lookupSymbol(candidate)) + if (!module.lookupSymbol(candidate)) { return candidate; + } } while (true); } @@ -185,8 +199,9 @@ static func::FuncOp createOutlinedHelper(ModuleOp module, SmallVector argTypes; argTypes.reserve(captures.size()); - for (Value capture : captures) + for (Value capture : captures) { argTypes.push_back(capture.getType()); + } OpBuilder moduleBuilder(module.getBodyRegion()); moduleBuilder.setInsertionPointToEnd(&module.getBodyRegion().front()); @@ -208,8 +223,9 @@ static func::FuncOp createOutlinedHelper(ModuleOp module, OpBuilder bodyBuilder = OpBuilder::atBlockEnd(entry); cloneExternalConstants(sectionOp, bodyBuilder, mapping); - for (Operation &op : sectionOp.getBody().front()) + for (Operation &op : sectionOp.getBody().front()) { bodyBuilder.clone(op, mapping); + } bodyBuilder.create(loc); return helper; @@ -234,16 +250,19 @@ static void replaceSectionWithLaunch(pto::SectionSimtOp sectionOp, static LogicalResult outlineSection(ModuleOp module, pto::SectionSimtOp sectionOp, unsigned &outlineIndex) { - if (failed(verifySectionCanBeOutlined(sectionOp))) + if (failed(verifySectionCanBeOutlined(sectionOp))) { return failure(); + } FailureOr maxThreads = getSimtThreadCount(sectionOp); - if (failed(maxThreads)) + if (failed(maxThreads)) { return failure(); + } SmallVector captures; - if (failed(collectCaptures(sectionOp, captures))) + if (failed(collectCaptures(sectionOp, captures))) { return failure(); + } func::FuncOp parentFunc = sectionOp->getParentOfType(); std::string helperName = diff --git a/lib/PTO/Transforms/PTOPlanMemory.cpp b/lib/PTO/Transforms/PTOPlanMemory.cpp index bdc7b8fdbe..e3aae82449 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.cpp +++ b/lib/PTO/Transforms/PTOPlanMemory.cpp @@ -65,26 +65,30 @@ struct LocalMemSpec { static std::optional getTileBufferFootprintBytes(TileBufType type) { ArrayRef shape = type.getShape(); unsigned elemBytes = getPTOStorageElemByteSize(type.getElementType()); - if (elemBytes == 0) + if (elemBytes == 0) { return std::nullopt; + } if (type.getCompactModeI32() != static_cast(pto::CompactMode::RowPlusOne)) { std::optional totalStaticSize = getStaticTotalSize(shape); - if (!totalStaticSize.has_value()) + if (!totalStaticSize.has_value()) { return std::nullopt; + } return totalStaticSize.value() * static_cast(elemBytes); } - if (shape.size() != 2 || llvm::is_contained(shape, ShapedType::kDynamic)) + if (shape.size() != 2 || llvm::is_contained(shape, ShapedType::kDynamic)) { return std::nullopt; + } bool rowMajor = type.getBLayoutValueI32() == static_cast(pto::BLayout::RowMajor); int64_t major = rowMajor ? shape[0] : shape[1]; int64_t minor = rowMajor ? shape[1] : shape[0]; - if (major == 0 || minor == 0) + if (major == 0 || minor == 0) { return 0; + } return ((major - 1) * (minor + 1) + minor) * static_cast(elemBytes); } @@ -95,11 +99,13 @@ static int64_t ceilDivBitsToBytes(int64_t bits) { static int64_t alignUpBytes(int64_t value, int64_t align) { int64_t safeAlign = std::max(align, 1); - if (safeAlign == 1) + if (safeAlign == 1) { return value; + } int64_t rem = value % safeAlign; - if (rem == 0) + if (rem == 0) { return value; + } return value + (safeAlign - rem); } @@ -130,19 +136,22 @@ static bool isIgnoredA5TmpOperandUse(OpOperand &use) { StringRef name = owner->getName().getStringRef(); if (auto dpsOp = dyn_cast(owner)) { - if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) + if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) { return false; + } } else if (auto dpsOp = dyn_cast(owner)) { - if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) + if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) { return false; + } } if (isNameIn(name, {"pto.trowargmax", "pto.trowargmin", "pto.trowmax", "pto.trowmin", "pto.trowsum", "pto.trowprod"})) return operandNo == 1; - if (name == "pto.txors") + if (name == "pto.txors") { return operandNo == 2; + } if (isNameIn(name, {"pto.tprelu", "pto.txor", "pto.tsels", "pto.trowexpand", "pto.tcolexpand", @@ -155,23 +164,27 @@ static bool isIgnoredA5TmpOperandUse(OpOperand &use) { "pto.tcolexpandmul", "pto.tcolexpandsub"})) return operandNo == 2; - if (name == "pto.tsel") + if (name == "pto.tsel") { return operandNo == 3; + } return false; } static bool isA5IgnoredTmpAlloc(pto::AllocTileOp allocTile) { - if (getTargetArch(allocTile.getOperation()) != PTOArch::A5) + if (getTargetArch(allocTile.getOperation()) != PTOArch::A5) { return false; + } Value value = allocTile.getResult(); - if (value.use_empty()) + if (value.use_empty()) { return false; + } for (OpOperand &use : value.getUses()) { - if (!isIgnoredA5TmpOperandUse(use)) + if (!isIgnoredA5TmpOperandUse(use)) { return false; + } } return true; } @@ -181,8 +194,9 @@ static void collectStableValueOrder(Region ®ion, DenseMap &stableValueKeys, SmallVectorImpl &seenValues) { auto recordValue = [&](Value value) { - if (stableValueKeys.find(value) != stableValueKeys.end()) + if (stableValueKeys.find(value) != stableValueKeys.end()) { return; + } std::string key; llvm::raw_string_ostream os(key); value.printAsOperand(os, asmState); @@ -191,14 +205,17 @@ static void collectStableValueOrder(Region ®ion, }; for (Block &block : region) { - for (BlockArgument blockArg : block.getArguments()) + for (BlockArgument blockArg : block.getArguments()) { recordValue(blockArg); + } for (Operation &op : block) { - for (Value result : op.getResults()) + for (Value result : op.getResults()) { recordValue(result); - for (Region &nestedRegion : op.getRegions()) + } + for (Region &nestedRegion : op.getRegions()) { collectStableValueOrder(nestedRegion, asmState, stableValueKeys, seenValues); + } } } } @@ -212,22 +229,25 @@ static StableValueOrderMap buildStableValueOrder(func::FuncOp func) { llvm::sort(seenValues, [&](Value lhs, Value rhs) { const std::string &lhsKey = stableValueKeys.find(lhs)->second; const std::string &rhsKey = stableValueKeys.find(rhs)->second; - if (lhsKey != rhsKey) + if (lhsKey != rhsKey) { return lhsKey < rhsKey; + } return isLessValue(lhs, rhs); }); StableValueOrderMap stableValueOrder; - for (auto [index, value] : llvm::enumerate(seenValues)) + for (auto [index, value] : llvm::enumerate(seenValues)) { stableValueOrder[value] = index; + } return stableValueOrder; } static uint32_t lookupStableValueOrder( Value value, const StableValueOrderMap &stableValueOrder) { auto it = stableValueOrder.find(value); - if (it != stableValueOrder.end()) + if (it != stableValueOrder.end()) { return it->second; + } return std::numeric_limits::max(); } @@ -237,8 +257,9 @@ static void sortValuesByStableOrder( llvm::sort(values, [&](Value lhs, Value rhs) { uint32_t lhsOrder = lookupStableValueOrder(lhs, stableValueOrder); uint32_t rhsOrder = lookupStableValueOrder(rhs, stableValueOrder); - if (lhsOrder != rhsOrder) + if (lhsOrder != rhsOrder) { return lhsOrder < rhsOrder; + } return isLessValue(lhs, rhs); }); } @@ -248,25 +269,31 @@ static SmallVector getScratchBuffersFromEffects(Operation *op, const StableValueOrderMap &stableValueOrder) { SmallVector scratchBuffers; auto memEffect = dyn_cast(op); - if (!memEffect) + if (!memEffect) { return scratchBuffers; + } SmallVector, kMemoryEffectReserveSize> effects; memEffect.getEffects(effects); for (const auto &effect : effects) { - if (!isa(effect.getEffect())) + if (!isa(effect.getEffect())) { continue; + } Value value = effect.getValue(); - if (!value) + if (!value) { continue; - if (!llvm::is_contained(op->getOperands(), value)) + } + if (!llvm::is_contained(op->getOperands(), value)) { continue; - if (llvm::is_contained(dpsInits, value)) + } + if (llvm::is_contained(dpsInits, value)) { continue; - if (!llvm::is_contained(scratchBuffers, value)) + } + if (!llvm::is_contained(scratchBuffers, value)) { scratchBuffers.push_back(value); + } } sortValuesByStableOrder(scratchBuffers, stableValueOrder); return scratchBuffers; @@ -277,21 +304,25 @@ getMemoryEffectBufferOperands(Operation *op, const StableValueOrderMap &stableValueOrder) { SmallVector buffers; auto memEffect = dyn_cast(op); - if (!memEffect) + if (!memEffect) { return buffers; + } SmallVector, kMemoryEffectReserveSize> effects; memEffect.getEffects(effects); for (const auto &effect : effects) { - if (!isa(effect.getEffect())) + if (!isa(effect.getEffect())) { continue; + } Value value = effect.getValue(); - if (!value || !GetBufferSpaceAttr(value)) + if (!value || !GetBufferSpaceAttr(value)) { continue; - if (!llvm::is_contained(buffers, value)) + } + if (!llvm::is_contained(buffers, value)) { buffers.push_back(value); + } } sortValuesByStableOrder(buffers, stableValueOrder); return buffers; @@ -305,8 +336,9 @@ getScratchConflictPairsFromEffects(Operation *op, ValueRange dpsInits, getScratchBuffersFromEffects(op, dpsInits, stableValueOrder); for (Value scratch : scratchBuffers) { for (Value dst : dpsInits) { - if (!scratch || !dst || scratch == dst) + if (!scratch || !dst || scratch == dst) { continue; + } conflictPairs.emplace_back(scratch, dst); } } @@ -336,14 +368,16 @@ static LogicalResult analyzeReserveBufferPlans(func::FuncOp funcOp, funcOp.walk( [&](ReserveBufferOp reserveOp) { reserveOps.push_back(reserveOp); }); - if (reserveOps.empty()) + if (reserveOps.empty()) { return success(); + } for (ReserveBufferOp reserveOp : reserveOps) { AddressSpace as = reserveOp.getLocation().getAddressSpace(); auto spec = getLocalMemSpec(reserveOp.getOperation(), as); - if (spec.capacityBits <= 0 || spec.alignBytes <= 0) + if (spec.capacityBits <= 0 || spec.alignBytes <= 0) { return reserveOp.emitOpError("unsupported reserve_buffer location"); + } int64_t capacityBytes = spec.capacityBits / kBitsPerByte; int64_t sizeBytes = reserveOp.getSize(); @@ -370,8 +404,9 @@ static LogicalResult analyzeReserveBufferPlans(func::FuncOp funcOp, // In manual mode, reserve_buffer.base is already fixed by the frontend or // an earlier stage. Only basic validation is needed here. auto baseAttr = reserveOp.getBaseAttr(); - if (!baseAttr) + if (!baseAttr) { return reserveOp.emitOpError("expects 'base' when 'auto' is false"); + } int64_t baseBytes = baseAttr.getInt(); if (baseBytes % spec.alignBytes != 0) { @@ -401,8 +436,9 @@ static LogicalResult assignAutoReserveBufferBases( const BufferInfo &bufferInfo = it.second; auto offsetsIt = buffer2Offsets.find(buffer); - if (offsetsIt == buffer2Offsets.end()) + if (offsetsIt == buffer2Offsets.end()) { continue; + } // Reserve-buffer allocation intentionally happens after normal MemPlan. // Reconstruct the already occupied byte ranges from the planned local @@ -437,12 +473,14 @@ static LogicalResult assignAutoReserveBufferBases( ranges.swap(merged); }; - for (auto &it : occupiedByAddressSpace) + for (auto &it : occupiedByAddressSpace) { normalizeRanges(it.second); + } for (ReserveBufferPlan &plan : plans) { - if (plan.mode != ReserveBufferMode::Auto || !plan.reserveOp) + if (plan.mode != ReserveBufferMode::Auto || !plan.reserveOp) { continue; + } SmallVector &occupied = occupiedByAddressSpace[plan.addressSpace]; @@ -453,8 +491,9 @@ static LogicalResult assignAutoReserveBufferBases( int64_t candidateBase = 0; for (const OccupiedByteRange &range : occupied) { candidateBase = alignUpBytes(candidateBase, plan.alignBytes); - if (candidateBase + plan.sizeBytes <= range.begin) + if (candidateBase + plan.sizeBytes <= range.begin) { break; + } candidateBase = std::max(candidateBase, range.end); } candidateBase = alignUpBytes(candidateBase, plan.alignBytes); @@ -538,8 +577,9 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { if (allocTileOp.getAddr()) { return WalkResult::advance(); } - if (isA5IgnoredTmpAlloc(allocTileOp)) + if (isA5IgnoredTmpAlloc(allocTileOp)) { return WalkResult::advance(); + } auto memorySpaceAttr = GetBufferSpaceAttr(allocTileOp.getResult()); if (!isLocalBuffer(memorySpaceAttr)) { allocTileOp.emitError("Alloc tile buffer not at local space"); @@ -562,8 +602,9 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { return WalkResult::advance(); } else if (isLocalMemPlan() && dyn_cast(op)) { auto allocMultiOp = cast(op); - if (allocMultiOp.getAddr()) + if (allocMultiOp.getAddr()) { return WalkResult::advance(); + } auto memorySpaceAttr = GetBufferSpaceAttr(allocMultiOp.getResult()); if (!isLocalBuffer(memorySpaceAttr)) { allocMultiOp.emitError("Alloc multi tile buffer not at local space"); @@ -700,8 +741,9 @@ void MemLivenessAnalysis::UpdateForOpBufferAlias(scf::ForOp forOp) { UpdateBufferAlias(forOp.getYieldedValues()[i], arg); } } - if (forOp->getResults().size() != forOp.getYieldedValues().size()) + if (forOp->getResults().size() != forOp.getYieldedValues().size()) { llvm::report_fatal_error("scf.for result/yield sizes are inconsistent"); + } for (auto [i, arg] : llvm::enumerate(forOp.getYieldedValues())) { // forOp result values alias region iter yielded values. UpdateBufferAlias(forOp->getResult(i), arg); @@ -723,8 +765,9 @@ void MemLivenessAnalysis::UpdateForOpInitArgsAlias(scf::ForOp forOp) { if (forOp.getInitArgs().empty()) { return; } - if (forOp.getInitArgs().size() != forOp.getRegionIterArgs().size()) + if (forOp.getInitArgs().size() != forOp.getRegionIterArgs().size()) { llvm::report_fatal_error("scf.for init/iter-arg sizes are inconsistent"); + } for (auto [i, arg] : llvm::enumerate(forOp.getInitArgs())) { // init args alias region iter args. UpdateBufferAlias(forOp.getRegionIterArgs()[i], arg); @@ -736,8 +779,9 @@ void MemLivenessAnalysis::UpdateIfOpBufferAlias(scf::IfOp ifOp, if (ifOp.getResults().empty()) { return; } - if (ifOp->getResults().size() != yieldOp->getOperands().size()) + if (ifOp->getResults().size() != yieldOp->getOperands().size()) { llvm::report_fatal_error("scf.if result/yield sizes are inconsistent"); + } for (auto [i, arg] : llvm::enumerate(yieldOp->getOperands())) { // Multiple buffers involved, requiring one-to-one correspondence. UpdateBufferAlias(ifOp->getResult(i), arg); @@ -762,8 +806,9 @@ void MemLivenessAnalysis::RecursiveIfOp(scf::IfOp ifOp, Liveness live) { void MemLivenessAnalysis::UpdateFusionRegionBufferAlias( pto::FusionRegionOp fusionRegion, pto::YieldOp yieldOp) { - if (fusionRegion.getResults().empty()) + if (fusionRegion.getResults().empty()) { return; + } if (fusionRegion->getResults().size() != yieldOp->getOperands().size()) { llvm::report_fatal_error( "pto.fusion_region result/yield sizes are inconsistent"); @@ -780,8 +825,9 @@ void MemLivenessAnalysis::RecursiveFusionRegionOp(pto::FusionRegionOp fusionRegi auto yieldOp = dyn_cast(fusionRegion.getBody().front().getTerminator()); - if (!yieldOp) + if (!yieldOp) { llvm::report_fatal_error("pto.fusion_region must terminate with pto.yield"); + } UpdateFusionRegionBufferAlias(fusionRegion, yieldOp); auto regionEnd = UpdateLinearOperation(fusionRegion.getOperation()); @@ -807,8 +853,9 @@ SmallVector MemLivenessAnalysis::GetLiveBuffersInLoop(scf::ForOp forOp, aliasBuffers.insert(operand); for (auto Buffer : aliasBuffers) { auto iter = buffer2status.find(Buffer); - if (iter != buffer2status.end()) + if (iter != buffer2status.end()) { allocBeforeLoopBuffers.push_back(Buffer); + } } } sortValuesByStableOrder(allocBeforeLoopBuffers, stableValueOrder); @@ -818,7 +865,6 @@ SmallVector MemLivenessAnalysis::GetLiveBuffersInLoop(scf::ForOp forOp, bool MemLivenessAnalysis::isSkippableOp(Operation *op) const { // Call-like ops are still modeled explicitly. Only pure terminators and // dim queries are skipped here. - // return isa(op); } @@ -871,8 +917,9 @@ SetVector MemLivenessAnalysis::Union(SetVector set1, } SetVector MemLivenessAnalysis::GetAliasBuffers(Value aliasBuffer) { - if (!aliasBuffer) + if (!aliasBuffer) { return {}; + } auto trueVar = buffer2AliasVec.find(aliasBuffer); if (trueVar != buffer2AliasVec.end()) { @@ -909,8 +956,9 @@ void MemLivenessAnalysis::UpdateOpGenInfo(OpInfo *opInfo, void MemLivenessAnalysis::UpdateOperandGenInfo(OpInfo *opInfo, Value operand) { auto iter_buffer = buffer2status.find(operand); - if (iter_buffer == buffer2status.end()) + if (iter_buffer == buffer2status.end()) { return; + } if (iter_buffer->second == BufferStatus::DEFFINED) { genKillMap[opInfo].gen.push_back(operand); buffer2status[iter_buffer->first] = BufferStatus::GENED; @@ -942,8 +990,9 @@ void MemLivenessAnalysis::UpdateOpKillInfo(OpInfo *opInfo, Value operand, aliasBuffers.insert(operand); for (Value aliasBuffer : aliasBuffers) { auto iterBuffer = buffer2status.find(aliasBuffer); - if (iterBuffer == buffer2status.end()) + if (iterBuffer == buffer2status.end()) { return; + } if (iterBuffer->second == BufferStatus::GENED && IsInSameBlock(iterBuffer->first.getDefiningOp(), opInfo->operation) && AllDeadAfter(opInfo->operation, aliasBuffers, live)) { @@ -974,24 +1023,29 @@ void MemLivenessAnalysis::RecordSemanticConflict(Value lhs, Value rhs) { rhsAliases.insert(rhs); auto appendUniquePair = [&](Value a, Value b) { - if (!a || !b || a == b) + if (!a || !b || a == b) { return; + } ValuePair pair = isLessValue(a, b) ? ValuePair(a, b) : ValuePair(b, a); - if (!llvm::is_contained(semanticConflictPairs, pair)) + if (!llvm::is_contained(semanticConflictPairs, pair)) { semanticConflictPairs.push_back(pair); + } }; - for (Value a : lhsAliases) - for (Value b : rhsAliases) + for (Value a : lhsAliases) { + for (Value b : rhsAliases) { appendUniquePair(a, b); + } + } } BufferInfo MemLivenessAnalysis::GenerateBufferInfo(Operation *op, Value operand) { auto memorySpaceAttr = GetBufferSpaceAttr(operand); if (isLocalMemPlan() && isLocalBuffer(memorySpaceAttr)) { - if (!memorySpaceAttr.has_value()) + if (!memorySpaceAttr.has_value()) { llvm::report_fatal_error("local buffer must have memory space"); + } return GetBufferInfo(op, operand, memorySpaceAttr.value().getAddressSpace()); } @@ -1017,9 +1071,10 @@ BufferInfo MemLivenessAnalysis::GetBufferInfo(Operation *op, Value operand, llvm_unreachable("local memory planner expects tile buffer roots"); } bufferInfo.bufferType = elementType; - if (!footprintBytes.has_value()) + if (!footprintBytes.has_value()) { llvm::report_fatal_error( "failed to obtain buffer static physical footprint"); + } bufferInfo.constBits = footprintBytes.value() * kBitsPerByte; return bufferInfo; } @@ -1042,8 +1097,9 @@ void MemLivenessAnalysis::GenerateBufferLife() { // Time given to buffer end. for (const Value &killBuffer : it->second.kill) { auto iter = buffer2Life.find(killBuffer); - if (iter == buffer2Life.end()) + if (iter == buffer2Life.end()) { llvm::report_fatal_error("buffer lifetime killed before generation"); + } iter->second->freeTime = scopeTime; } scopeTime++; @@ -1062,8 +1118,9 @@ StorageEntry::GetBufferLifeByValue(const Value v) const { } bool MemPlan::IsReusePTOOp(Operation *op) const { - if (restrictInplaceAsISA) + if (restrictInplaceAsISA) { return false; + } // not in ISA but confirmed with hardware developers: // elementwise ops with the same shape and the same bitwidth operands can also @@ -1078,8 +1135,9 @@ SmallVector MemPlan::GenerateInplaceList() { inplacePairList.end()); for (auto &operationSeq : linearOperation) { auto it = genKillMap.find(operationSeq.get()); - if (it == genKillMap.end()) + if (it == genKillMap.end()) { continue; + } if (hasTouchOp[operationSeq->operation]) { continue; } @@ -1091,16 +1149,18 @@ SmallVector MemPlan::GenerateInplaceList() { for (const Value &genBuffer : genBuffers) { auto genBufferIter = bufferInfos.find(genBuffer); - if (genBufferIter == bufferInfos.end()) + if (genBufferIter == bufferInfos.end()) { llvm::report_fatal_error("gen buffer missing from buffer info map"); + } if (genBufferIter->second.ignoreInplace) { continue; } for (const Value &killBuffer : killBuffers) { auto killBufferIter = bufferInfos.find(killBuffer); - if (killBufferIter == bufferInfos.end()) + if (killBufferIter == bufferInfos.end()) { llvm::report_fatal_error("kill buffer missing from buffer info map"); + } if (killBufferIter->second.ignoreInplace) { continue; } @@ -1122,8 +1182,9 @@ SmallVector MemPlan::GenerateInplaceList() { } void MemPlan::EmitPlanMemoryFailureInfo() { - if (failApplyBufferInfo.empty()) + if (failApplyBufferInfo.empty()) { return; + } for (auto &iter : failApplyBufferInfo) { AddressSpace space = iter.first; func_.emitError() << stringifyEnum(space) << " overflow, requires " @@ -1165,8 +1226,9 @@ bool MemPlan::RecordOverflowIfAny() { bool MemPlan::HasSemanticConflict(const StorageEntry *entry, const BufferLifeVec &bufferLives) const { - if (!entry || semanticConflictPairs.empty() || bufferLives.empty()) + if (!entry || semanticConflictPairs.empty() || bufferLives.empty()) { return false; + } auto containsPair = [&](Value lhs, Value rhs) { ValuePair pair = isLessValue(lhs, rhs) ? ValuePair(lhs, rhs) @@ -1176,13 +1238,16 @@ bool MemPlan::HasSemanticConflict(const StorageEntry *entry, for (Value entryBuffer : entry->inplaceBuffers) { for (const auto &life : bufferLives) { - if (!life) + if (!life) { continue; + } Value otherBuffer = life->buffer; - if (!otherBuffer || entryBuffer == otherBuffer) + if (!otherBuffer || entryBuffer == otherBuffer) { continue; - if (containsPair(entryBuffer, otherBuffer)) + } + if (containsPair(entryBuffer, otherBuffer)) { return true; + } } } return false; @@ -1257,8 +1322,9 @@ void MemPlan::GenerateStorageEntry() { // create new storage entry. for (auto &operation : linearOperation) { auto it = genKillMap.find(operation.get()); - if (it == genKillMap.end()) + if (it == genKillMap.end()) { continue; + } SmallVector genBuffers(it->second.gen.begin(), it->second.gen.end()); sortValuesByStableOrder(genBuffers, stableValueOrder); for (const Value &genBuffer : genBuffers) { @@ -1286,8 +1352,9 @@ void MemPlan::GenerateStorageEntry() { void MemPlan::PrintSuccessfulAllocatedMaxBits() { auto it = memscope2rootStorageEntry.find(pto::AddressSpace::VEC); if (it != memscope2rootStorageEntry.end()) { - if (!it->second) + if (!it->second) { llvm::report_fatal_error("missing root storage entry for VEC scope"); + } uint64_t ubAllocBits = it->second->alignedConstBits + it->second->bitsOffset; for (auto& child : it->second->mergedChildren) { ubAllocBits = std::max(ubAllocBits, child->bitsOffset + child->alignedConstBits); @@ -1298,12 +1365,15 @@ void MemPlan::PrintSuccessfulAllocatedMaxBits() { } void MemPlan::ValidateParameters(std::unique_ptr &e) const { - if (!e->bufInfo->operation) + if (!e->bufInfo->operation) { llvm::report_fatal_error("storage entry missing defining operation"); - if (e->bufInfo->constBits < 0U) + } + if (e->bufInfo->constBits < 0U) { llvm::report_fatal_error("storage entry has invalid memory size"); - if (e->bufferLifeVec.empty()) + } + if (e->bufferLifeVec.empty()) { llvm::report_fatal_error("storage entry missing lifetime information"); + } } void MemPlan::UpdateBuffer2Offsets() { @@ -1313,16 +1383,18 @@ void MemPlan::UpdateBuffer2Offsets() { // skip the sibling offsets would be appended in StorageEntryVec order // rather than slot order, breaking the runtime contract that // `buffer2Offsets[buffer][k]` is slot k's physical offset. - if (e->isMultiBufferSlot) + if (e->isMultiBufferSlot) { continue; + } for (Value &buffer : e->inplaceBuffers) { buffer2Offsets[buffer].push_back( (e->bitsOffset + kBitsToByte - 1) / kBitsToByte); // Multi-buffer primary: append sibling offsets in slot order so the // final offsets list is [slot0, slot1, ..., slotN-1]. for (auto *sibling : e->relationOtherBuffers) { - if (!sibling) + if (!sibling) { continue; + } buffer2Offsets[buffer].push_back( (sibling->bitsOffset + kBitsToByte - 1) / kBitsToByte); } @@ -1359,8 +1431,9 @@ void MemPlan::MergeInplaceSE() { // already same storageEntry, no need to inplace. continue; } - if (genSE == nullptr || killSE == nullptr) + if (genSE == nullptr || killSE == nullptr) { llvm::report_fatal_error("invalid storage entry during inplace merge"); + } BufferLifeVec mergedBufferLifeVec; mergedBufferLifeVec.insert(mergedBufferLifeVec.end(), genSE->bufferLifeVec.begin(), @@ -1463,8 +1536,9 @@ void MemPlan::ExpandMultiBufferStorageEntry() { size_t size = StorageEntryVec.size(); for (size_t i = 0; i < size; i++) { auto *primary = StorageEntryVec[i].get(); - if (primary->multiBufferNum <= 1) + if (primary->multiBufferNum <= 1) { continue; + } uint32_t n = primary->multiBufferNum; for (uint32_t slot = 1; slot < n; ++slot) { auto entry = std::make_unique(); @@ -1477,8 +1551,9 @@ void MemPlan::ExpandMultiBufferStorageEntry() { primary->relationOtherBuffers.push_back(entry.get()); StorageEntryVec.push_back(std::move(entry)); } - if (!primary->relationOtherBuffers.empty()) + if (!primary->relationOtherBuffers.empty()) { primary->relationPongEntry = primary->relationOtherBuffers.front(); + } } } @@ -1487,8 +1562,9 @@ bool MemPlan::IsEnoughForBuffersNoReuse(StorageEntry *rootStorageEntry, size_t alignUnit) { auto iter = bufferScope2RequiredSize.find(rootStorageEntry->bufInfo->bufferScope); - if (iter == bufferScope2RequiredSize.end()) + if (iter == bufferScope2RequiredSize.end()) { llvm::report_fatal_error("missing required-size entry for buffer scope"); + } if (iter->second < restBufferSize) { // Even when the scope fits without reuse (no peak to save), honor // largest-first placement so the option means the same thing on both paths: @@ -1595,16 +1671,18 @@ PlanStatus MemPlan::PlanMemAddressOfWholeLocalBuffer() { size_t maxBits = bufferSpaceInfo.second; if (rootStorageEntry->mergedChildren.empty()) { PlanStatus status = PlanSingleLocalBuffer(rootStorageEntry, align, maxBits); - if (status != PlanStatus::PLAN_SUCCESS) + if (status != PlanStatus::PLAN_SUCCESS) { return status; + } continue; } if (IsEnoughForBuffersNoReuse(rootStorageEntry, maxBits, align)) { continue; } PlanStatus status = PlanReusableLocalBuffer(rootStorageEntry, align, maxBits); - if (status != PlanStatus::PLAN_SUCCESS) + if (status != PlanStatus::PLAN_SUCCESS) { return status; + } } planStatus = PlanStatus::PLAN_SUCCESS; return planStatus; @@ -1666,8 +1744,9 @@ PlanStatus MemPlan::PlanReusableLocalBuffer(StorageEntry *rootStorageEntry, return status; } } - if (si.childIdx >= childrenNum) + if (si.childIdx >= childrenNum) { break; + } curEntry = rootStorageEntry->mergedChildren[si.childIdx]; } return PlanStatus::PLAN_SUCCESS; @@ -1932,8 +2011,9 @@ MemPlan::GetBufferParentLoop(const SmallVector &buffers) { llvm::SmallSet parentLoopVec; for (auto buffer : buffers) { if (!buffer.getDefiningOp()) { - if (!isa(buffer.getParentBlock()->getParentOp())) + if (!isa(buffer.getParentBlock()->getParentOp())) { llvm::report_fatal_error("expected loop-carried block argument"); + } // Init args and region iter arg are inplace, ignore Region Iter Arg // without DefineOp. continue; @@ -2035,8 +2115,9 @@ void MemPlan::SpecAllocRelationPongEntry(MemBoundList &outline, PlanRecHis &his, if (e->multiBufferNum == kDoubleBufferCount && e->relationPongEntry) { pongStorageEntry = e->relationPongEntry; } - if (!pongStorageEntry) + if (!pongStorageEntry) { llvm::report_fatal_error("pong storage entry not found"); + } UpdateOutline(outline, his, pongStorageEntry, OutlineSectionInfo(start, end, size, true), SPEC_LEVEL_1); return; @@ -2048,8 +2129,9 @@ bool MemPlan::IsBufferLifeVecConflict(PlanRecord &r, uint64_t offset, const StorageEntry *e) const { if ((r.firstMemBound->offset + r.allExtent > offset) && (r.firstMemBound->offset < offset + e->alignedConstBits)) { - if (HasSemanticConflict(e, r.firstMemBound->bufferLifeVec)) + if (HasSemanticConflict(e, r.firstMemBound->bufferLifeVec)) { return true; + } DenseMap intersection = GetOverlapBufferLife(r.entry->bufferLifeVec, e->bufferLifeVec); return !intersection.empty(); @@ -2266,8 +2348,9 @@ bool MemPlan::IsSamePlanAsLastRollBack(uint64_t allocOffset, int curChildIdx, inline bool MemPlan::VerifyConflictStage0(StorageEntry *e, const std::shared_ptr &last) { - if (HasSemanticConflict(e, last->bufferLifeVec)) + if (HasSemanticConflict(e, last->bufferLifeVec)) { return true; + } // level_0: offset = 0, offset means life distance DenseMap intersection = GetOverlapBufferLife(e->bufferLifeVec, last->bufferLifeVec); @@ -2368,8 +2451,9 @@ void MemPlan::ReportAllocatedEntryDebugInfo(StorageEntry *rootStorageEntry) { LDBG("\n"); } size_t num = allocatedEntry.size() - 1; - if (rootStorageEntry->mergedChildren.size() <= num) + if (rootStorageEntry->mergedChildren.size() <= num) { llvm::report_fatal_error("missing failed storage entry"); + } const StorageEntry *failedSe = rootStorageEntry->mergedChildren[num]; printRecord(failedSe); LDBG("alloc fail,because exceed bound of memory \n" @@ -2522,16 +2606,19 @@ class LegacyAllocTileOpAddPlannedAddressPattern LogicalResult matchAndRewrite(pto::AllocTileOp op, PatternRewriter &rewriter) const override { - if (op.getAddr()) + if (op.getAddr()) { return failure(); + } auto tileType = dyn_cast(op.getResult().getType()); - if (!tileType) + if (!tileType) { return failure(); + } auto it = buffer2Offsets.find(op.getResult()); - if (it == buffer2Offsets.end() || it->second.empty()) + if (it == buffer2Offsets.end() || it->second.empty()) { return failure(); + } if (it->second.size() != 1) { return rewriter.notifyMatchFailure( @@ -2545,8 +2632,9 @@ class LegacyAllocTileOpAddPlannedAddressPattern op.getValidRow() ? op.getValidRow() : Value(), op.getValidCol() ? op.getValidCol() : Value()); for (NamedAttribute attr : op->getAttrs()) { - if (attr.getName().getValue() == "operandSegmentSizes") + if (attr.getName().getValue() == "operandSegmentSizes") { continue; + } planned->setAttr(attr.getName(), attr.getValue()); } @@ -2569,11 +2657,13 @@ class LegacyAllocMultiTileOpAddPlannedAddressesPattern LogicalResult matchAndRewrite(pto::AllocMultiTileOp op, PatternRewriter &rewriter) const override { - if (op.getAddr() || op->hasAttr(pto::kPtoMultiBufferAddrsAttrName)) + if (op.getAddr() || op->hasAttr(pto::kPtoMultiBufferAddrsAttrName)) { return failure(); + } auto it = buffer2Offsets.find(op.getResult()); - if (it == buffer2Offsets.end() || it->second.empty()) + if (it == buffer2Offsets.end() || it->second.empty()) { return failure(); + } if (it->second.size() != op.getResult().getType().getCount()) { return rewriter.notifyMatchFailure( op, "planned address count does not match multi_tile_buf count"); @@ -2581,8 +2671,9 @@ class LegacyAllocMultiTileOpAddPlannedAddressesPattern SmallVector addrs; addrs.reserve(it->second.size()); - for (uint64_t offset : it->second) + for (uint64_t offset : it->second) { addrs.push_back(static_cast(offset)); + } rewriter.modifyOpInPlace(op, [&] { op->setAttr(pto::kPtoMultiBufferAddrsAttrName, rewriter.getDenseI64ArrayAttr(addrs)); @@ -2599,8 +2690,9 @@ static FailureOr parseLegacyMemPlanMode(func::FuncOp func, if (memMode.equals_insensitive("local") || memMode.equals_insensitive("local-mem-plan")) return MemPlanMode::LOCAL_MEM_PLAN; - if (memMode.equals_insensitive("global-work-space-plan")) + if (memMode.equals_insensitive("global-work-space-plan")) { return MemPlanMode::GLOBAL_WORKSPACE_PLAN; + } func.emitError("unsupported mem-mode '") << memMode << "'; only 'local' is supported by the PTOAS pipeline"; return failure(); @@ -2635,14 +2727,16 @@ void PlanMemoryPass::runOnOperation() { // TileOp helpers only contain compute code and deliberately do not own // alloc_tile/reserve_buffer lifetimes. All other functions, including // ordinary functions in backend child modules, must be planned. - if (!funcOp->hasAttr("pto.tileop.helper")) + if (!funcOp->hasAttr("pto.tileop.helper")) { funcs.push_back(funcOp); + } }); for (func::FuncOp funcOp : funcs) { auto parsedMode = parseLegacyMemPlanMode(funcOp, this->memMode); - if (failed(parsedMode)) + if (failed(parsedMode)) { return signalPassFailure(); + } MemPlanMode mode = *parsedMode; ReserveBufferPlans reservePlans; if (mode == MemPlanMode::LOCAL_MEM_PLAN && @@ -2651,8 +2745,9 @@ void PlanMemoryPass::runOnOperation() { } if (mode == MemPlanMode::LOCAL_MEM_PLAN) { for (ReserveBufferPlan &reservePlan : reservePlans) { - if (reservePlan.mode != ReserveBufferMode::Manual) + if (reservePlan.mode != ReserveBufferMode::Manual) { continue; + } reservePlan.reserveOp.emitOpError( "pto.reserve_buffer with explicit 'base' (auto = false) is not " "supported in PlanMemory; use --pto-level=level3 or set auto = true"); @@ -2700,18 +2795,22 @@ void PlanMemoryPass::runOnOperation() { bool hasUnplannedAllocTile = false; funcOp.walk([&](pto::AllocTileOp op) { - if (op.getAddr()) + if (op.getAddr()) { return; - if (op->use_empty()) + } + if (op->use_empty()) { return; - if (isA5IgnoredTmpAlloc(op)) + } + if (isA5IgnoredTmpAlloc(op)) { return; + } op.emitError( "PTOPlanMemory failed to assign an address to pto.alloc_tile"); hasUnplannedAllocTile = true; }); - if (hasUnplannedAllocTile) + if (hasUnplannedAllocTile) { return signalPassFailure(); + } } } diff --git a/lib/PTO/Transforms/PTOPlanMemory.h b/lib/PTO/Transforms/PTOPlanMemory.h index 65bcdc6317..71f8a186b2 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.h +++ b/lib/PTO/Transforms/PTOPlanMemory.h @@ -37,7 +37,7 @@ struct ValueComparator { using StableValueOrderMap = DenseMap; /// Various states when collecting gen-kill. -enum BufferStatus { UNDEFFINED = 0, DEFFINED, GENED, KILLED }; +enum class BufferStatus { UNDEFFINED = 0, DEFFINED, GENED, KILLED }; /// Pair of inplace Value. using ValuePair = std::pair; @@ -48,7 +48,7 @@ enum class MemPlanMode { }; /// Result status after plan memory. -enum PlanStatus { +enum class PlanStatus { PLAN_SUCCESS = 0, RESTART_NEW_PLAN, CONTINUE_PLAN, @@ -116,7 +116,7 @@ struct GenKillEntry { struct BufferLife { BufferLife(Value buffer, int64_t start, int64_t end) : buffer(buffer), allocTime(start), freeTime(end) {} - BufferLife(Value buffer) : buffer(buffer) {} + explicit BufferLife(Value buffer) : buffer(buffer) {} /// buffer value. Value buffer; /// the buffer allocate time. @@ -805,7 +805,6 @@ class MemPlan { /// The device's SCALING storage size int scalingSpaceSize{0}; - }; } // namespace pto } // namespace mlir diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 64ea9f9c5f..4fdc343dea 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -95,21 +95,24 @@ struct ReuseGroup { }; static uint64_t alignUp(uint64_t value, uint64_t align) { - if (align == 0) + if (align == 0) { return value; + } return ((value + align - 1) / align) * align; } static std::optional getBufferAddressSpace(Type type) { if (auto tileType = dyn_cast(type)) { if (auto attr = - dyn_cast_or_null(tileType.getMemorySpace())) + dyn_cast_or_null(tileType.getMemorySpace())) { return attr.getAddressSpace(); + } return std::nullopt; } - if (auto multiType = dyn_cast(type)) + if (auto multiType = dyn_cast(type)) { return getBufferAddressSpace(multiType.getSlotType()); + } return std::nullopt; } @@ -133,19 +136,23 @@ static bool isIgnoredA5TmpOperandUse(OpOperand &use) { StringRef name = owner->getName().getStringRef(); if (auto dpsOp = dyn_cast(owner)) { - if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) + if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) { return false; + } } else if (auto dpsOp = dyn_cast(owner)) { - if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) + if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) { return false; + } } if (isNameIn(name, {"pto.trowargmax", "pto.trowargmin", "pto.trowmax", - "pto.trowmin", "pto.trowsum", "pto.trowprod"})) + "pto.trowmin", "pto.trowsum", "pto.trowprod"})) { return operandNo == 1; + } - if (name == "pto.txors") + if (name == "pto.txors") { return operandNo == 2; + } if (isNameIn(name, {"pto.tprelu", "pto.txor", "pto.tsels", "pto.trowexpand", "pto.tcolexpand", @@ -158,23 +165,27 @@ static bool isIgnoredA5TmpOperandUse(OpOperand &use) { "pto.tcolexpandmul", "pto.tcolexpandsub"})) return operandNo == 2; - if (name == "pto.tsel") + if (name == "pto.tsel") { return operandNo == 3; + } return false; } static bool isA5IgnoredTmpAlloc(pto::AllocTileOp allocTile) { - if (getTargetArch(allocTile.getOperation()) != PTOArch::A5) + if (getTargetArch(allocTile.getOperation()) != PTOArch::A5) { return false; + } Value value = allocTile.getResult(); - if (value.use_empty()) + if (value.use_empty()) { return false; + } for (OpOperand &use : value.getUses()) { - if (!isIgnoredA5TmpOperandUse(use)) + if (!isIgnoredA5TmpOperandUse(use)) { return false; + } } return true; } @@ -182,18 +193,18 @@ static bool isA5IgnoredTmpAlloc(pto::AllocTileOp allocTile) { static MemSpec getMemSpec(PTOArch arch, AddressSpace space) { switch (space) { case AddressSpace::VEC: - return {arch == PTOArch::A5 ? 253952ull : 196608ull, 256}; + return {arch == PTOArch::A5 ? 253952ULL : 196608ULL, 256}; case AddressSpace::MAT: - return {524288ull, 256}; + return {524288ULL, 256}; case AddressSpace::LEFT: case AddressSpace::RIGHT: - return {65536ull, 4096}; + return {65536ULL, 4096}; case AddressSpace::ACC: - return {arch == PTOArch::A5 ? 262144ull : 131072ull, 4096}; + return {arch == PTOArch::A5 ? 262144ULL : 131072ULL, 4096}; case AddressSpace::BIAS: - return {65536ull, 256}; + return {65536ULL, 256}; case AddressSpace::SCALING: - return {arch == PTOArch::A5 ? 253952ull : 196608ull, 256}; + return {arch == PTOArch::A5 ? 253952ULL : 196608ULL, 256}; case AddressSpace::GM: case AddressSpace::Zero: break; @@ -205,50 +216,58 @@ static FailureOr computeStaticBufferBytes(Value value) { auto computeTileBytes = [](TileBufType type) -> FailureOr { ArrayRef shape = type.getShape(); uint64_t elemBytes = getPTOStorageElemByteSize(type.getElementType()); - if (elemBytes == 0) + if (elemBytes == 0) { return failure(); + } if (type.getCompactModeI32() == static_cast(CompactMode::RowPlusOne)) { if (shape.size() != 2 || - llvm::is_contained(shape, ShapedType::kDynamic)) + llvm::is_contained(shape, ShapedType::kDynamic)) { return failure(); + } bool rowMajor = type.getBLayoutValueI32() == static_cast(BLayout::RowMajor); uint64_t major = static_cast(rowMajor ? shape[0] : shape[1]); uint64_t minor = static_cast(rowMajor ? shape[1] : shape[0]); - if (major == 0 || minor == 0) + if (major == 0 || minor == 0) { return uint64_t{0}; + } return ((major - 1) * (minor + 1) + minor) * elemBytes; } uint64_t numel = 1; for (int64_t dim : shape) { - if (dim == ShapedType::kDynamic) + if (dim == ShapedType::kDynamic) { return failure(); + } numel *= static_cast(dim); } return numel * elemBytes; }; - if (auto tileType = dyn_cast(value.getType())) + if (auto tileType = dyn_cast(value.getType())) { return computeTileBytes(tileType); - if (auto multiType = dyn_cast(value.getType())) + } + if (auto multiType = dyn_cast(value.getType())) { return computeTileBytes(multiType.getSlotType()); + } return failure(); } static void appendUniqueRoot(RootList &roots, Value root) { - if (llvm::is_contained(roots, root)) + if (llvm::is_contained(roots, root)) { return; + } roots.push_back(root); } static RootList unionRoots(const RootList &lhs, const RootList &rhs) { RootList result = lhs; - for (Value root : rhs) + for (Value root : rhs) { appendUniqueRoot(result, root); + } return result; } @@ -314,39 +333,48 @@ static SmallVector getWrittenNonDpsOperands(Operation *op, ValueRange dpsInits) { SmallVector scratchOperands; auto memEffect = dyn_cast(op); - if (!memEffect) + if (!memEffect) { return scratchOperands; + } SmallVector, 8> effects; memEffect.getEffects(effects); for (const auto &effect : effects) { - if (!isa(effect.getEffect())) + if (!isa(effect.getEffect())) { continue; + } Value value = effect.getValue(); - if (!value) + if (!value) { continue; - if (!llvm::is_contained(op->getOperands(), value)) + } + if (!llvm::is_contained(op->getOperands(), value)) { continue; - if (llvm::is_contained(dpsInits, value)) + } + if (llvm::is_contained(dpsInits, value)) { continue; - if (!llvm::is_contained(scratchOperands, value)) + } + if (!llvm::is_contained(scratchOperands, value)) { scratchOperands.push_back(value); + } } return scratchOperands; } static bool hasReadEffectOnValue(Operation *op, Value value) { auto memEffect = dyn_cast(op); - if (!memEffect) + if (!memEffect) { return false; + } SmallVector, 8> effects; memEffect.getEffects(effects); for (const auto &effect : effects) { - if (!isa(effect.getEffect())) + if (!isa(effect.getEffect())) { continue; - if (effect.getValue() == value) + } + if (effect.getValue() == value) { return true; + } } return false; } @@ -357,10 +385,12 @@ static bool isInsideLoop(Operation *op) { } static ValueRange getDpsInits(Operation *op) { - if (auto dpsOp = dyn_cast(op)) + if (auto dpsOp = dyn_cast(op)) { return dpsOp.getDpsInits(); - if (auto dpsOp = dyn_cast(op)) + } + if (auto dpsOp = dyn_cast(op)) { return dpsOp.getDpsInits(); + } return ValueRange(); } @@ -382,30 +412,35 @@ struct PlannerAnalysis { RootList getRoots(Value value) const { auto it = valueToRoots.find(value); - if (it == valueToRoots.end()) + if (it == valueToRoots.end()) { return {}; + } return it->second; } void setRoots(Value value, const RootList &roots) { - if (roots.empty()) + if (roots.empty()) { return; + } valueToRoots[value] = roots; } void mergeRoots(Value value, const RootList &roots) { - if (roots.empty()) + if (roots.empty()) { return; + } setRoots(value, unionRoots(getRoots(value), roots)); } void addRoot(Value value, Operation *defOp) { - if (rootIndexByValue.count(value)) + if (rootIndexByValue.count(value)) { return; + } auto space = getBufferAddressSpace(value.getType()); - if (!isPlannableLocalSpace(space)) + if (!isPlannableLocalSpace(space)) { return; + } auto bytesOr = computeStaticBufferBytes(value); if (mlir::failed(bytesOr)) { @@ -416,8 +451,9 @@ struct PlannerAnalysis { } uint64_t slotCount = 1; - if (auto multiType = dyn_cast(value.getType())) + if (auto multiType = dyn_cast(value.getType())) { slotCount = multiType.getCount(); + } else if (auto attr = defOp->getAttrOfType(pto::kPtoMultiBufferAttrName)) slotCount = attr.getValue().getZExtValue(); @@ -441,13 +477,16 @@ struct PlannerAnalysis { void markUse(Value value, unsigned index) { for (Value root : getRoots(value)) { auto found = rootIndexByValue.find(root); - if (found == rootIndexByValue.end()) + if (found == rootIndexByValue.end()) { continue; + } RootInfo &info = roots[found->second]; - if (!info.hasWriter) + if (!info.hasWriter) { info.hasUseBeforeFirstWrite = true; - if (index < info.allocIndex) + } + if (index < info.allocIndex) { info.allocIndex = index; + } info.freeIndex = std::max(info.freeIndex, index); } } @@ -455,12 +494,14 @@ struct PlannerAnalysis { void markWrite(Value value, unsigned index, bool pureOverwrite) { for (Value root : getRoots(value)) { auto found = rootIndexByValue.find(root); - if (found == rootIndexByValue.end()) + if (found == rootIndexByValue.end()) { continue; + } RootInfo &info = roots[found->second]; if (!info.hasWriter) { - if (pureOverwrite && !info.hasUseBeforeFirstWrite) + if (pureOverwrite && !info.hasUseBeforeFirstWrite) { info.allocIndex = index; + } info.hasWriter = true; } info.freeIndex = std::max(info.freeIndex, index); @@ -468,33 +509,40 @@ struct PlannerAnalysis { } void addForbidAlias(Value a, Value b) { - if (a == b) + if (a == b) { return; - if (!rootIndexByValue.count(a) || !rootIndexByValue.count(b)) + } + if (!rootIndexByValue.count(a) || !rootIndexByValue.count(b)) { return; + } appendUniqueRoot(facts.forbidAlias[a], b); appendUniqueRoot(facts.forbidAlias[b], a); } bool hasForbidAlias(Value a, Value b) const { - if (a == b) + if (a == b) { return false; + } auto it = facts.forbidAlias.find(a); - if (it == facts.forbidAlias.end()) + if (it == facts.forbidAlias.end()) { return false; + } return llvm::is_contained(it->second, b); } void addForbidAliasBetweenRoots(const RootList &lhsRoots, const RootList &rhsRoots) { - for (Value lhs : lhsRoots) - for (Value rhs : rhsRoots) + for (Value lhs : lhsRoots) { + for (Value rhs : rhsRoots) { addForbidAlias(lhs, rhs); + } + } } void markRoots(DenseSet &set, const RootList &roots) { - for (Value root : roots) + for (Value root : roots) { set.insert(root); + } } bool rootsContain(const DenseSet &set, const RootList &roots) const { @@ -518,53 +566,64 @@ struct PlannerAnalysis { } void propagateSplitTpopDerived(Value result, ValueRange sources) { - if (!result) + if (!result) { return; + } if (llvm::any_of(sources, [&](Value source) { return isSplitTpopDerived(source); - })) + })) { splitTpopDerivedValues.insert(result); + } } void propagateSplitTpopDerivedFromRoots(Value result, Value source) { - if (!result || !source) + if (!result || !source) { return; - if (isSplitTpopDerived(source)) + } + if (isSplitTpopDerived(source)) { splitTpopDerivedValues.insert(result); + } } void propagateSplitTpopDerivedFromOperands(Operation *op) { - if (!operandsContainSplitTpopDerived(op)) + if (!operandsContainSplitTpopDerived(op)) { return; - for (Value result : op->getResults()) + } + for (Value result : op->getResults()) { splitTpopDerivedValues.insert(result); + } } bool isRootDefinedInRegion(Value root, Region ®ion) const { auto it = rootIndexByValue.find(root); - if (it == rootIndexByValue.end()) + if (it == rootIndexByValue.end()) { return false; + } Operation *defOp = roots[it->second].defOp; for (Operation *cur = defOp; cur; cur = cur->getParentOp()) { - if (cur->getParentRegion() == ®ion) + if (cur->getParentRegion() == ®ion) { return true; + } } return false; } void addBranchExclusivePair(Value lhs, Value rhs) { - if (lhs == rhs) + if (lhs == rhs) { return; - if (!rootIndexByValue.count(lhs) || !rootIndexByValue.count(rhs)) + } + if (!rootIndexByValue.count(lhs) || !rootIndexByValue.count(rhs)) { return; + } appendUniqueRoot(facts.branchExclusiveRoots[lhs], rhs); appendUniqueRoot(facts.branchExclusiveRoots[rhs], lhs); } void recordIfBranchExclusivity(scf::IfOp ifOp) { - if (ifOp.getNumResults() == 0) + if (ifOp.getNumResults() == 0) { return; + } auto thenYield = cast(ifOp.thenBlock()->getTerminator()); auto elseYield = cast(ifOp.elseBlock()->getTerminator()); @@ -572,31 +631,40 @@ struct PlannerAnalysis { llvm::zip(thenYield.getResults(), elseYield.getResults())) { RootList thenLocalRoots; RootList elseLocalRoots; - for (Value root : getRoots(thenVal)) - if (isRootDefinedInRegion(root, ifOp.getThenRegion())) + for (Value root : getRoots(thenVal)) { + if (isRootDefinedInRegion(root, ifOp.getThenRegion())) { appendUniqueRoot(thenLocalRoots, root); - for (Value root : getRoots(elseVal)) - if (isRootDefinedInRegion(root, ifOp.getElseRegion())) + } + } + for (Value root : getRoots(elseVal)) { + if (isRootDefinedInRegion(root, ifOp.getElseRegion())) { appendUniqueRoot(elseLocalRoots, root); + } + } - for (Value thenRoot : thenLocalRoots) - for (Value elseRoot : elseLocalRoots) + for (Value thenRoot : thenLocalRoots) { + for (Value elseRoot : elseLocalRoots) { addBranchExclusivePair(thenRoot, elseRoot); + } + } } } void recordDpsScratchConflicts(Operation *op, ValueRange dpsInits) { RootList outputRoots; - for (Value init : dpsInits) + for (Value init : dpsInits) { outputRoots = unionRoots(outputRoots, getRoots(init)); - if (outputRoots.empty()) + } + if (outputRoots.empty()) { return; + } for (Value scratch : getWrittenNonDpsOperands(op, dpsInits)) { addForbidAliasBetweenRoots(getRoots(scratch), outputRoots); for (Value operand : op->getOperands()) { - if (operand == scratch || llvm::is_contained(dpsInits, operand)) + if (operand == scratch || llvm::is_contained(dpsInits, operand)) { continue; + } addForbidAliasBetweenRoots(getRoots(scratch), getRoots(operand)); } } @@ -604,34 +672,40 @@ struct PlannerAnalysis { void recordInplacePolicyConflicts(Operation *op, ValueRange dpsInits) { RootList outputRoots; - for (Value init : dpsInits) + for (Value init : dpsInits) { outputRoots = unionRoots(outputRoots, getRoots(init)); - if (outputRoots.empty()) + } + if (outputRoots.empty()) { return; + } InplacePolicy policy = getInplacePolicy(op); if (policy.notInplaceSafe) { for (Value operand : op->getOperands()) { - if (llvm::is_contained(dpsInits, operand)) + if (llvm::is_contained(dpsInits, operand)) { continue; + } addForbidAliasBetweenRoots(getRoots(operand), outputRoots); } } for (unsigned operandIndex : policy.forbidOutputAliasOperands) { - if (operandIndex >= op->getNumOperands()) + if (operandIndex >= op->getNumOperands()) { continue; + } Value operand = op->getOperand(operandIndex); - if (llvm::is_contained(dpsInits, operand)) + if (llvm::is_contained(dpsInits, operand)) { continue; + } addForbidAliasBetweenRoots(getRoots(operand), outputRoots); } } void recordDpsInplaceConflicts(Operation *op) { auto dpsOp = dyn_cast(op); - if (!dpsOp) + if (!dpsOp) { return; + } ValueRange dpsInits = dpsOp.getDpsInits(); recordDpsScratchConflicts(op, dpsInits); @@ -639,35 +713,42 @@ struct PlannerAnalysis { } void recordLoadDerivedRoots(Operation *op, ValueRange dpsInits) { - if (!isa(op)) + if (!isa(op)) { return; + } - for (Value init : dpsInits) + for (Value init : dpsInits) { markRoots(facts.loadDerivedRoots, getRoots(init)); + } } void recordTpopConsumerRoots(Operation *op, ValueRange dpsInits, unsigned opIndex) { - if (!facts.targetHazardEnabled) + if (!facts.targetHazardEnabled) { return; - if (!operandsContainSplitTpopDerived(op) || !operandsContainLoadDerivedRoot(op)) + } + if (!operandsContainSplitTpopDerived(op) || !operandsContainLoadDerivedRoot(op)) { return; + } RootList outputRoots; - for (Value init : dpsInits) + for (Value init : dpsInits) { outputRoots = unionRoots(outputRoots, getRoots(init)); + } markRoots(facts.tpopConsumerRoots, outputRoots); for (Value root : outputRoots) { SmallVector &indices = facts.tpopConsumerWriteIndices[root]; - if (!llvm::is_contained(indices, opIndex)) + if (!llvm::is_contained(indices, opIndex)) { indices.push_back(opIndex); + } } } void recordDpsTargetHazardFacts(Operation *op, unsigned opIndex) { auto dpsOp = dyn_cast(op); - if (!dpsOp) + if (!dpsOp) { return; + } ValueRange dpsInits = dpsOp.getDpsInits(); recordLoadDerivedRoots(op, dpsInits); @@ -675,34 +756,41 @@ struct PlannerAnalysis { } void appendAccessRoots(RootList &dst, Value value) { - for (Value root : getRoots(value)) + for (Value root : getRoots(value)) { appendUniqueRoot(dst, root); + } } void recordRootAccessStats(ArrayRef accessRoots, PIPE pipe, bool isWrite, bool inLoop) { for (Value root : accessRoots) { auto found = rootIndexByValue.find(root); - if (found == rootIndexByValue.end()) + if (found == rootIndexByValue.end()) { continue; + } RootInfo &info = roots[found->second]; ++info.accessCount; - if (isWrite) + if (isWrite) { ++info.writeAccessCount; - if (inLoop) + } + if (inLoop) { ++info.loopAccessCount; - if (pipe == PIPE::PIPE_V) + } + if (pipe == PIPE::PIPE_V) { ++info.pipeVAccessCount; - if (pipe == PIPE::PIPE_MTE2 || pipe == PIPE::PIPE_MTE3) + } + if (pipe == PIPE::PIPE_MTE2 || pipe == PIPE::PIPE_MTE3) { ++info.mteAccessCount; + } } } void recordOpAccess(Operation *op, ValueRange dpsInits) { auto pipeOp = dyn_cast(op); auto memEffect = dyn_cast(op); - if (!pipeOp || !memEffect) + if (!pipeOp || !memEffect) { return; + } OpAccess access; access.op = op; @@ -714,19 +802,23 @@ struct PlannerAnalysis { memEffect.getEffects(effects); for (const auto &effect : effects) { Value value = effect.getValue(); - if (!value) + if (!value) { continue; - if (isa(effect.getEffect())) + } + if (isa(effect.getEffect())) { appendAccessRoots(access.reads, value); - if (isa(effect.getEffect())) + } + if (isa(effect.getEffect())) { appendAccessRoots(access.writes, value); + } } // Some PTO tile ops expose DPS destinations through the project-specific // interface even when a future op forgets to model MemoryEffects. Keep the // performance side table conservative by treating DPS inits as writes. - for (Value init : dpsInits) + for (Value init : dpsInits) { appendAccessRoots(access.writes, init); + } if (!access.reads.empty() || !access.writes.empty()) { recordRootAccessStats(access.reads, access.pipe, /*isWrite=*/false, @@ -738,18 +830,21 @@ struct PlannerAnalysis { } void recordSplitTpopDerivedValue(Operation *op) { - if (!facts.targetHazardEnabled) + if (!facts.targetHazardEnabled) { return; + } if (auto pop = dyn_cast(op)) { - if (pop.getSplit() != 0) + if (pop.getSplit() != 0) { splitTpopDerivedValues.insert(pop.getTile()); + } return; } if (auto pop = dyn_cast(op)) { - if (pop.getSplit() != 0) + if (pop.getSplit() != 0) { splitTpopDerivedValues.insert(pop.getTile()); + } return; } @@ -757,8 +852,9 @@ struct PlannerAnalysis { } void seedForIterArgAliases(scf::ForOp forOp) { - if (forOp.getRegion().empty()) + if (forOp.getRegion().empty()) { return; + } Block &body = forOp.getRegion().front(); for (auto [iterArg, initArg] : llvm::zip(body.getArguments().drop_front(1), forOp.getInitArgs())) { @@ -777,8 +873,9 @@ struct PlannerAnalysis { SmallVector successors; branchOp.getSuccessorRegions(RegionBranchPoint::parent(), successors); for (RegionSuccessor successor : successors) { - if (successor.isParent()) + if (successor.isParent()) { continue; + } mapRegionBranchValues(branchOp.getEntrySuccessorOperands(successor), successor.getSuccessorInputs()); } @@ -786,21 +883,24 @@ struct PlannerAnalysis { void finalizeRegionBranchAliasesFromRegion(RegionBranchOpInterface branchOp, Region ®ion) { - if (region.empty()) + if (region.empty()) { return; + } SmallVector successors; branchOp.getSuccessorRegions(region, successors); for (RegionSuccessor successor : successors) { ValueRange destinations = successor.getSuccessorInputs(); - if (destinations.empty()) + if (destinations.empty()) { continue; + } for (Block &block : region) { auto terminator = dyn_cast(block.getTerminator()); - if (!terminator) + if (!terminator) { continue; + } mapRegionBranchValues(terminator.getSuccessorOperands(successor), destinations); } @@ -808,8 +908,9 @@ struct PlannerAnalysis { } void finalizeRegionBranchAliases(RegionBranchOpInterface branchOp) { - for (Region ®ion : branchOp->getRegions()) + for (Region ®ion : branchOp->getRegions()) { finalizeRegionBranchAliasesFromRegion(branchOp, region); + } } void finalizeForLoopLiveness(scf::ForOp forOp, unsigned loopStartIndex, @@ -818,10 +919,12 @@ struct PlannerAnalysis { // across the back-edge, even if its last use appears before a loop-local // allocation in the linear walk. for (RootInfo &info : roots) { - if (isRootDefinedInRegion(info.root, forOp.getRegion())) + if (isRootDefinedInRegion(info.root, forOp.getRegion())) { continue; - if (info.freeIndex >= loopStartIndex) + } + if (info.freeIndex >= loopStartIndex) { info.freeIndex = std::max(info.freeIndex, loopEndIndex); + } } auto yieldOp = cast(forOp.getBody()->getTerminator()); @@ -838,8 +941,9 @@ struct PlannerAnalysis { // do not apply branch-exclusivity from one iteration across the back-edge. for (Value root : loopCarriedRoots) { auto found = rootIndexByValue.find(root); - if (found == rootIndexByValue.end()) + if (found == rootIndexByValue.end()) { continue; + } facts.loopCarriedRoots.insert(root); RootInfo &info = roots[found->second]; info.allocIndex = std::min(info.allocIndex, loopStartIndex); @@ -856,26 +960,31 @@ struct PlannerAnalysis { // cycle so a later linear operation in one region cannot reuse storage that // a future iteration still reads. RootList loopCarriedRoots; - for (Value init : whileOp.getInits()) + for (Value init : whileOp.getInits()) { loopCarriedRoots = unionRoots(loopCarriedRoots, getRoots(init)); - for (Value result : whileOp.getResults()) + } + for (Value result : whileOp.getResults()) { loopCarriedRoots = unionRoots(loopCarriedRoots, getRoots(result)); + } for (Region ®ion : whileOp->getRegions()) { for (Block &block : region) { - for (BlockArgument arg : block.getArguments()) + for (BlockArgument arg : block.getArguments()) { loopCarriedRoots = unionRoots(loopCarriedRoots, getRoots(arg)); + } if (auto terminator = dyn_cast( block.getTerminator())) { - for (Value operand : terminator->getOperands()) + for (Value operand : terminator->getOperands()) { loopCarriedRoots = unionRoots(loopCarriedRoots, getRoots(operand)); + } } } } for (Value root : loopCarriedRoots) { auto found = rootIndexByValue.find(root); - if (found == rootIndexByValue.end()) + if (found == rootIndexByValue.end()) { continue; + } facts.loopCarriedRoots.insert(root); RootInfo &info = roots[found->second]; info.allocIndex = std::min(info.allocIndex, loopStartIndex); @@ -893,11 +1002,13 @@ struct PlannerAnalysis { if (auto allocTile = dyn_cast(op)) { if (!allocTile.getAddr()) { - if (isA5IgnoredTmpAlloc(allocTile)) + if (isA5IgnoredTmpAlloc(allocTile)) { continue; + } addRoot(allocTile.getResult(), op); - if (failed) + if (failed) { return; + } auto found = rootIndexByValue.find(allocTile.getResult()); if (found != rootIndexByValue.end()) { roots[found->second].allocIndex = index; @@ -907,8 +1018,9 @@ struct PlannerAnalysis { } else if (auto allocMulti = dyn_cast(op)) { if (!allocMulti.getAddr()) { addRoot(allocMulti.getResult(), op); - if (failed) + if (failed) { return; + } auto found = rootIndexByValue.find(allocMulti.getResult()); if (found != rootIndexByValue.end()) { roots[found->second].allocIndex = index; @@ -946,38 +1058,44 @@ struct PlannerAnalysis { ValueRange dpsInits = getDpsInits(op); for (Value operand : op->getOperands()) { - if (llvm::is_contained(dpsInits, operand)) + if (llvm::is_contained(dpsInits, operand)) { continue; + } markUse(operand, index); } for (Value init : dpsInits) { bool readsOldValue = hasReadEffectOnValue(op, init); - if (readsOldValue) + if (readsOldValue) { markUse(init, index); + } markWrite(init, index, /*pureOverwrite=*/!readsOldValue); } if (auto branchOp = dyn_cast(op)) { for (Region &nested : op->getRegions()) { walkRegion(nested); - if (failed) + if (failed) { return; + } finalizeRegionBranchAliasesFromRegion(branchOp, nested); } } else { for (Region &nested : op->getRegions()) { walkRegion(nested); - if (failed) + if (failed) { return; + } } } - if (auto branchOp = dyn_cast(op)) + if (auto branchOp = dyn_cast(op)) { finalizeRegionBranchAliases(branchOp); + } if (auto ifOp = dyn_cast(op)) { - if (ifOp.getNumResults() != 0) + if (ifOp.getNumResults() != 0) { recordIfBranchExclusivity(ifOp); + } } else if (auto forOp = dyn_cast(op)) { unsigned loopEndIndex = linearOps.empty() ? index : linearOps.size() - 1; @@ -990,8 +1108,9 @@ struct PlannerAnalysis { auto yieldOp = cast( fusionRegion.getBody().front().getTerminator()); for (auto [result, yielded] : - llvm::zip(fusionRegion.getResults(), yieldOp.getOperands())) + llvm::zip(fusionRegion.getResults(), yieldOp.getOperands())) { setRoots(result, getRoots(yielded)); + } } recordSplitTpopDerivedValue(op); @@ -1009,18 +1128,21 @@ static bool lifetimesStrictlyOverlap(const RootInfo &lhs, const RootInfo &rhs) { static bool areBranchExclusive(Value lhs, Value rhs, const ConflictFacts &facts) { if (facts.loopCarriedRoots.contains(lhs) || - facts.loopCarriedRoots.contains(rhs)) + facts.loopCarriedRoots.contains(rhs)) { return false; + } auto it = facts.branchExclusiveRoots.find(lhs); - if (it == facts.branchExclusiveRoots.end()) + if (it == facts.branchExclusiveRoots.end()) { return false; + } return llvm::is_contained(it->second, rhs); } static bool gateLifetimeAndPhi(const RootInfo &lhs, const RootInfo &rhs, const ConflictFacts &facts) { - if (!lifetimesStrictlyOverlap(lhs, rhs)) + if (!lifetimesStrictlyOverlap(lhs, rhs)) { return true; + } // Gate 5 is intentionally embedded in Gate 1: branch-local roots yielded from // opposite scf.if branches are mutually exclusive at runtime, so their @@ -1030,16 +1152,19 @@ static bool gateLifetimeAndPhi(const RootInfo &lhs, const RootInfo &rhs, static bool hasTargetHazard(const RootInfo &input, const RootInfo &writer, const ConflictFacts &facts) { - if (!facts.targetHazardEnabled) + if (!facts.targetHazardEnabled) { return false; + } if (!facts.loadDerivedRoots.contains(input.root) || - !facts.tpopConsumerRoots.contains(writer.root)) + !facts.tpopConsumerRoots.contains(writer.root)) { return false; + } auto found = facts.tpopConsumerWriteIndices.find(writer.root); - if (found == facts.tpopConsumerWriteIndices.end()) + if (found == facts.tpopConsumerWriteIndices.end()) { return false; + } return llvm::is_contained(found->second, input.freeIndex); } @@ -1064,11 +1189,13 @@ static bool canShare(const RootInfo &lhs, const RootInfo &rhs, static bool canJoinReuseGroup(const RootInfo &info, const ReuseGroup &group, const PlannerAnalysis &analysis) { - if (info.space != group.space) + if (info.space != group.space) { return false; + } for (const RootInfo *member : group.members) { - if (!canShare(info, *member, analysis)) + if (!canShare(info, *member, analysis)) { return false; + } } return true; } @@ -1116,8 +1243,9 @@ static uint64_t loopWeightedPenalty(uint64_t penalty, const OpAccess &lhs, static uint64_t getRootPairReuseCost(Value lhsRoot, Value rhsRoot, const ConflictFacts &facts) { - if (lhsRoot == rhsRoot) + if (lhsRoot == rhsRoot) { return 0; + } constexpr unsigned kPipeVLookahead = 1; constexpr unsigned kMteLookahead = 1; @@ -1126,14 +1254,17 @@ static uint64_t getRootPairReuseCost(Value lhsRoot, Value rhsRoot, uint64_t cost = 0; for (const OpAccess &lhs : facts.opAccesses) { - if (!accessesRoot(lhs, lhsRoot) && !accessesRoot(lhs, rhsRoot)) + if (!accessesRoot(lhs, lhsRoot) && !accessesRoot(lhs, rhsRoot)) { continue; + } for (const OpAccess &rhs : facts.opAccesses) { - if (lhs.opIndex >= rhs.opIndex) + if (lhs.opIndex >= rhs.opIndex) { continue; - if (!accessesRoot(rhs, lhsRoot) && !accessesRoot(rhs, rhsRoot)) + } + if (!accessesRoot(rhs, lhsRoot) && !accessesRoot(rhs, rhsRoot)) { continue; + } if (lhs.pipe == PIPE::PIPE_V && rhs.pipe == PIPE::PIPE_V && isNearNeighbor(lhs, rhs, kPipeVLookahead) && @@ -1148,8 +1279,9 @@ static uint64_t getRootPairReuseCost(Value lhsRoot, Value rhsRoot, containsRoot(rhs.writes, rhsRoot)) || (containsRoot(lhs.reads, rhsRoot) && containsRoot(rhs.writes, lhsRoot)); - if (storeSourceThenLoadDst) + if (storeSourceThenLoadDst) { cost += loopWeightedPenalty(kMte3ToMte2Penalty, lhs, rhs); + } } } } @@ -1160,16 +1292,18 @@ static uint64_t getGroupReuseCost(const RootInfo &info, const ReuseGroup &group, const ConflictFacts &facts) { uint64_t cost = 0; - for (const RootInfo *member : group.members) + for (const RootInfo *member : group.members) { cost += getRootPairReuseCost(info.root, member->root, facts); + } return cost; } static bool isHotRoot(const RootInfo &info) { if (info.loopAccessCount > 0 && (info.accessCount >= 2 || info.pipeVAccessCount > 0 || - info.mteAccessCount > 0)) + info.mteAccessCount > 0)) { return true; + } // Some PTODSL kernels express repeated work through task/block parallelism // rather than an explicit scf.for in PTO IR. Repeated local accesses on @@ -1189,8 +1323,9 @@ static uint64_t getRootHotness(const RootInfo &info) { static uint64_t getHotClusterReuseCost(const RootInfo &info, const ReuseGroup &group) { - if (!isHotRoot(info)) + if (!isHotRoot(info)) { return 0; + } constexpr uint64_t kHotClusterPenalty = 6; constexpr uint64_t kLoopHotClusterPenalty = 12; @@ -1200,22 +1335,25 @@ static uint64_t getHotClusterReuseCost(const RootInfo &info, uint64_t cost = 0; uint64_t infoHotness = getRootHotness(info); for (const RootInfo *member : group.members) { - if (!isHotRoot(*member)) + if (!isHotRoot(*member)) { continue; + } uint64_t pairHotness = std::min(infoHotness, getRootHotness(*member)); cost += kHotClusterPenalty + pairHotness; - if (info.loopAccessCount > 0 && member->loopAccessCount > 0) + if (info.loopAccessCount > 0 && member->loopAccessCount > 0) { cost += kLoopHotClusterPenalty; + } // Exact co-location is the strongest possible same-bank signal. The // planner does not materialize final byte intervals yet, so only model the // bank pattern that is known for a reuse candidate: the new root would // start at the same offset as every member in this group. if (info.alignmentBytes <= kBankConflictModuloBytes && - member->alignmentBytes <= kBankConflictModuloBytes) + member->alignmentBytes <= kBankConflictModuloBytes) { cost += kBankConflictPenalty; + } } return cost; } @@ -1252,8 +1390,9 @@ static ReuseGroup *chooseReuseGroupByCost( bool bestFits = false; for (auto [index, group] : llvm::enumerate(groups)) { - if (!canJoinReuseGroup(info, group, analysis)) + if (!canJoinReuseGroup(info, group, analysis)) { continue; + } uint64_t projectedBytes = getProjectedPackedBytes(groups, info, static_cast(index)); @@ -1284,27 +1423,31 @@ static ReuseGroup *chooseReuseGroupByCost( // when local capacity is available. uint64_t freshCost = groups.empty() ? 0 : 1; unsigned freshOrder = groups.size(); - if (!bestGroup) + if (!bestGroup) { return nullptr; + } // The cost model is a performance hint, not a correctness gate. When local // memory is already tight, do not let a fresh address outrank a legal reuse // group; future roots may still need the remaining tail bytes. - if (bestFits && freshFits && freshSlack < pressureReserve) + if (bestFits && freshFits && freshSlack < pressureReserve) { return bestGroup; + } auto isBetter = [](bool lhsFits, uint64_t lhsCost, uint64_t lhsBytes, unsigned lhsOrder, bool rhsFits, uint64_t rhsCost, uint64_t rhsBytes, unsigned rhsOrder) { - if (lhsFits != rhsFits) + if (lhsFits != rhsFits) { return lhsFits; + } return std::tie(lhsCost, lhsBytes, lhsOrder) < std::tie(rhsCost, rhsBytes, rhsOrder); }; if (isBetter(freshFits, freshCost, freshProjectedBytes, freshOrder, bestFits, - bestCost, bestProjectedBytes, bestOrder)) + bestCost, bestProjectedBytes, bestOrder)) { return nullptr; + } return bestGroup; } @@ -1312,8 +1455,9 @@ static SmallVector buildSlotOffsets(uint64_t base, uint64_t slotBytes, uint64_t slotCount) { SmallVector offsets; offsets.reserve(slotCount); - for (uint64_t slot = 0; slot < slotCount; ++slot) + for (uint64_t slot = 0; slot < slotCount; ++slot) { offsets.push_back(base + slot * slotBytes); + } return offsets; } @@ -1338,13 +1482,15 @@ static FailureOr planReserveBufferBase( uint64_t cursor = 0; for (const auto &interval : merged) { cursor = alignUp(cursor, spec.alignmentBytes); - if (cursor + sizeBytes <= interval.first) + if (cursor + sizeBytes <= interval.first) { return cursor; + } cursor = std::max(cursor, interval.second); } cursor = alignUp(cursor, spec.alignmentBytes); - if (cursor + sizeBytes > spec.capacityBytes) + if (cursor + sizeBytes > spec.capacityBytes) { return failure(); + } occupied.push_back({cursor, cursor + sizeBytes}); return cursor; } @@ -1353,13 +1499,15 @@ static LogicalResult validateManualReserveBufferBase(pto::ReserveBufferOp reserveOp, const MemSpec &spec) { auto baseAttr = reserveOp.getBaseAttr(); - if (!baseAttr) + if (!baseAttr) { return reserveOp.emitError("expects 'base' when 'auto' is false"); + } int64_t signedBase = baseAttr.getInt(); - if (signedBase < 0) + if (signedBase < 0) { return reserveOp.emitError( "expects 'base' to be non-negative when present"); + } uint64_t base = static_cast(signedBase); if (base % spec.alignmentBytes != 0) { @@ -1393,8 +1541,9 @@ validateReserveBufferForPlanMemory(pto::ReserveBufferOp reserveOp, return success(); } - if (mlir::failed(validateManualReserveBufferBase(reserveOp, spec))) + if (mlir::failed(validateManualReserveBufferBase(reserveOp, spec))) { return failure(); + } return reserveOp.emitError( "pto.reserve_buffer with explicit 'base' (auto = false) is not " @@ -1412,16 +1561,19 @@ class AllocTileOpAddPlannedAddressPattern LogicalResult matchAndRewrite(pto::AllocTileOp op, PatternRewriter &rewriter) const override { - if (op.getAddr()) + if (op.getAddr()) { return failure(); + } auto tileType = dyn_cast(op.getResult().getType()); - if (!tileType) + if (!tileType) { return failure(); + } auto it = buffer2Offsets.find(op.getResult()); - if (it == buffer2Offsets.end() || it->second.empty()) + if (it == buffer2Offsets.end() || it->second.empty()) { return failure(); + } if (it->second.size() != 1) { return rewriter.notifyMatchFailure( @@ -1435,8 +1587,9 @@ class AllocTileOpAddPlannedAddressPattern op.getValidRow() ? op.getValidRow() : Value(), op.getValidCol() ? op.getValidCol() : Value()); for (NamedAttribute attr : op->getAttrs()) { - if (attr.getName().getValue() == "operandSegmentSizes") + if (attr.getName().getValue() == "operandSegmentSizes") { continue; + } planned->setAttr(attr.getName(), attr.getValue()); } @@ -1459,12 +1612,14 @@ class AllocMultiTileOpAddPlannedAddressesPattern LogicalResult matchAndRewrite(pto::AllocMultiTileOp op, PatternRewriter &rewriter) const override { - if (op.getAddr() || op->hasAttr(pto::kPtoMultiBufferAddrsAttrName)) + if (op.getAddr() || op->hasAttr(pto::kPtoMultiBufferAddrsAttrName)) { return failure(); + } auto it = buffer2Offsets.find(op.getResult()); - if (it == buffer2Offsets.end() || it->second.empty()) + if (it == buffer2Offsets.end() || it->second.empty()) { return failure(); + } if (it->second.size() != op.getResult().getType().getCount()) { return rewriter.notifyMatchFailure( op, "planned address count does not match multi_tile_buf count"); @@ -1472,8 +1627,9 @@ class AllocMultiTileOpAddPlannedAddressesPattern SmallVector addrs; addrs.reserve(it->second.size()); - for (uint64_t offset : it->second) + for (uint64_t offset : it->second) { addrs.push_back(static_cast(offset)); + } rewriter.modifyOpInPlace(op, [&] { op->setAttr(pto::kPtoMultiBufferAddrsAttrName, rewriter.getDenseI64ArrayAttr(addrs)); @@ -1492,8 +1648,9 @@ static LogicalResult materializePlannedOffsets( buffer2Offsets); patterns.add( patterns.getContext(), buffer2Offsets); - if (mlir::failed(applyPatternsGreedily(func, std::move(patterns)))) + if (mlir::failed(applyPatternsGreedily(func, std::move(patterns)))) { return failure(); + } return success(); } @@ -1516,8 +1673,9 @@ LogicalResult mlir::pto::runModernPlanMemory(func::FuncOp func, DenseMap> buffer2Offsets; llvm::MapVector> rootsBySpace; - for (RootInfo &info : analysis.roots) + for (RootInfo &info : analysis.roots) { rootsBySpace[info.space].push_back(&info); + } for (auto &entry : rootsBySpace) { AddressSpace space = entry.first; @@ -1526,10 +1684,12 @@ LogicalResult mlir::pto::runModernPlanMemory(func::FuncOp func, bool sizeFirstForSpace = orderBySize && !isCubeLocalSpace(space); llvm::stable_sort(roots, [&](const RootInfo *lhs, const RootInfo *rhs) { - if (sizeFirstForSpace && lhs->totalBytes != rhs->totalBytes) + if (sizeFirstForSpace && lhs->totalBytes != rhs->totalBytes) { return lhs->totalBytes > rhs->totalBytes; - if (lhs->allocIndex != rhs->allocIndex) + } + if (lhs->allocIndex != rhs->allocIndex) { return lhs->allocIndex < rhs->allocIndex; + } return lhs->stableOrder < rhs->stableOrder; }); @@ -1577,15 +1737,17 @@ LogicalResult mlir::pto::runModernPlanMemory(func::FuncOp func, return failure(); } - for (RootInfo *info : roots) + for (RootInfo *info : roots) { buffer2Offsets[info->root] = info->offsets; + } } DenseMap>> occupiedBySpace; for (const RootInfo &info : analysis.roots) { - if (info.offsets.empty()) + if (info.offsets.empty()) { continue; + } occupiedBySpace[info.space].push_back( {info.offsets.front(), info.offsets.front() + info.totalBytes}); } @@ -1621,17 +1783,21 @@ LogicalResult mlir::pto::runModernPlanMemory(func::FuncOp func, bool hasUnplannedAllocTile = false; func.walk([&](pto::AllocTileOp op) { - if (op.getAddr()) + if (op.getAddr()) { return; - if (op->use_empty()) + } + if (op->use_empty()) { return; - if (isA5IgnoredTmpAlloc(op)) + } + if (isA5IgnoredTmpAlloc(op)) { return; + } op.emitError("PTOPlanMemory failed to assign an address to pto.alloc_tile"); hasUnplannedAllocTile = true; }); - if (hasUnplannedAllocTile) + if (hasUnplannedAllocTile) { return failure(); + } return success(); } @@ -1664,8 +1830,9 @@ struct PlanMemoryModernPass // `pto.tileop.helper` identifies compute-only helpers, not whether a // child module needs memory planning. Skip those helpers themselves, // while planning every ordinary function in nested backend modules. - if (!funcOp->hasAttr("pto.tileop.helper")) + if (!funcOp->hasAttr("pto.tileop.helper")) { funcs.push_back(funcOp); + } }); for (func::FuncOp funcOp : funcs) { diff --git a/lib/PTO/Transforms/PTORematerializeFixpipeVectorQuant.cpp b/lib/PTO/Transforms/PTORematerializeFixpipeVectorQuant.cpp index 814e4e7567..4271b05e7a 100644 --- a/lib/PTO/Transforms/PTORematerializeFixpipeVectorQuant.cpp +++ b/lib/PTO/Transforms/PTORematerializeFixpipeVectorQuant.cpp @@ -53,8 +53,9 @@ struct PTORematerializeFixpipeVectorQuantPass auto processBlock = [&](auto &&self, Block &block) -> LogicalResult { llvm::DenseMap activeVectorById; SmallVector originalOps; - for (Operation &op : block) + for (Operation &op : block) { originalOps.push_back(&op); + } for (Operation *op : originalOps) { if (auto setQuantVector = dyn_cast(op)) { @@ -91,8 +92,9 @@ struct PTORematerializeFixpipeVectorQuantPass } } - for (Operation *op : eraseList) + for (Operation *op : eraseList) { op->erase(); + } } }; diff --git a/lib/PTO/Transforms/PTORemoveIdentityTMov.cpp b/lib/PTO/Transforms/PTORemoveIdentityTMov.cpp index 330b52820b..058da2dfdc 100644 --- a/lib/PTO/Transforms/PTORemoveIdentityTMov.cpp +++ b/lib/PTO/Transforms/PTORemoveIdentityTMov.cpp @@ -49,42 +49,48 @@ static std::optional tryEvalIntegerLikeConstant(Value value); static std::optional evalSignedCast(Value input, Type resultType) { std::optional inputValue = tryEvalIntegerLikeConstant(input); std::optional resultWidth = getIntegerLikeBitWidth(resultType); - if (!inputValue || !resultWidth) + if (!inputValue || !resultWidth) { return std::nullopt; + } return inputValue->sextOrTrunc(*resultWidth); } static std::optional evalUnsignedCast(Value input, Type resultType) { std::optional inputValue = tryEvalIntegerLikeConstant(input); std::optional resultWidth = getIntegerLikeBitWidth(resultType); - if (!inputValue || !resultWidth) + if (!inputValue || !resultWidth) { return std::nullopt; + } return inputValue->zextOrTrunc(*resultWidth); } static std::optional evalTruncCast(Value input, Type resultType) { std::optional inputValue = tryEvalIntegerLikeConstant(input); std::optional resultWidth = getIntegerLikeBitWidth(resultType); - if (!inputValue || !resultWidth || *resultWidth > inputValue->getBitWidth()) + if (!inputValue || !resultWidth || *resultWidth > inputValue->getBitWidth()) { return std::nullopt; + } return inputValue->trunc(*resultWidth); } static std::optional tryEvalIntegerLikeConstant(Value value) { - if (!value) + if (!value) { return std::nullopt; + } APInt apInt; if (matchPattern(value, m_ConstantInt(&apInt))) { std::optional width = getIntegerLikeBitWidth(value.getType()); - if (!width) + if (!width) { return std::nullopt; + } return apInt.sextOrTrunc(*width); } Operation *defOp = value.getDefiningOp(); - if (!defOp) + if (!defOp) { return std::nullopt; + } if (auto castOp = dyn_cast(defOp)) return evalSignedCast(castOp.getIn(), castOp.getType()); @@ -102,8 +108,9 @@ static std::optional tryEvalIntegerLikeConstant(Value value) { static std::optional tryEvalI64Constant(Value value) { std::optional apInt = tryEvalIntegerLikeConstant(value); - if (!apInt || apInt->getBitWidth() > 64) + if (!apInt || apInt->getBitWidth() > 64) { return std::nullopt; + } return apInt->getSExtValue(); } @@ -125,22 +132,26 @@ static bool isDeadDstTMov(TMovOp op) { static const BaseMemInfo * getSingleMemInfo(const Buffer2MemInfoMap &buffer2MemInfoMap, Value value) { auto it = buffer2MemInfoMap.find(value); - if (it == buffer2MemInfoMap.end() || it->second.size() != 1) + if (it == buffer2MemInfoMap.end() || it->second.size() != 1) { return nullptr; + } return it->second.front().get(); } static std::optional tryGetConcreteRootAddress(const BaseMemInfo *info) { - if (!info) + if (!info) { return std::nullopt; + } - if (auto direct = tryEvalI64Constant(info->rootBuffer)) + if (auto direct = tryEvalI64Constant(info->rootBuffer)) { return direct; + } Operation *defOp = info->rootBuffer.getDefiningOp(); - if (!defOp) + if (!defOp) { return std::nullopt; + } if (auto alloc = dyn_cast(defOp)) return tryEvalI64Constant(alloc.getAddr()); @@ -158,20 +169,23 @@ static bool isStaticallyAddressableValue(Value value) { constexpr int kMaxDepth = 32; while (value && depth++ < kMaxDepth) { Operation *defOp = value.getDefiningOp(); - if (!defOp) + if (!defOp) { return false; + } if (auto subView = dyn_cast(defOp)) { if (hasDynamicStaticList(subView.getStaticOffsets()) || hasDynamicStaticList(subView.getStaticSizes()) || - hasDynamicStaticList(subView.getStaticStrides())) + hasDynamicStaticList(subView.getStaticStrides())) { return false; + } value = subView.getSource(); continue; } - if (isa(defOp)) + if (isa(defOp)) { return false; + } if (auto cast = dyn_cast(defOp)) { value = cast.getSource(); @@ -186,8 +200,9 @@ static bool isStaticallyAddressableValue(Value value) { continue; } if (auto view = dyn_cast(defOp)) { - if (view.getByteShift()) + if (view.getByteShift()) { return false; + } value = view.getSource(); continue; } @@ -200,27 +215,35 @@ static bool isStaticallyAddressableValue(Value value) { static bool hasExactSameAddressRange(const BaseMemInfo *srcInfo, const BaseMemInfo *dstInfo) { - if (!srcInfo || !dstInfo) + if (!srcInfo || !dstInfo) { return false; - if (srcInfo->scope != dstInfo->scope) + } + if (srcInfo->scope != dstInfo->scope) { return false; - if (srcInfo->allocateSize == 0 || dstInfo->allocateSize == 0) + } + if (srcInfo->allocateSize == 0 || dstInfo->allocateSize == 0) { return false; - if (srcInfo->allocateSize != dstInfo->allocateSize) + } + if (srcInfo->allocateSize != dstInfo->allocateSize) { return false; - if (srcInfo->baseAddresses.empty() || dstInfo->baseAddresses.empty()) + } + if (srcInfo->baseAddresses.empty() || dstInfo->baseAddresses.empty()) { return false; - if (srcInfo->baseAddresses != dstInfo->baseAddresses) + } + if (srcInfo->baseAddresses != dstInfo->baseAddresses) { return false; + } return true; } static bool hasSameConcreteAddressRange(const BaseMemInfo *srcInfo, const BaseMemInfo *dstInfo) { - if (!hasExactSameAddressRange(srcInfo, dstInfo)) + if (!hasExactSameAddressRange(srcInfo, dstInfo)) { return false; - if (srcInfo->rootBuffer == dstInfo->rootBuffer) + } + if (srcInfo->rootBuffer == dstInfo->rootBuffer) { return true; + } auto srcRootAddr = tryGetConcreteRootAddress(srcInfo); auto dstRootAddr = tryGetConcreteRootAddress(dstInfo); return srcRootAddr && dstRootAddr && *srcRootAddr == *dstRootAddr; @@ -228,8 +251,9 @@ static bool hasSameConcreteAddressRange(const BaseMemInfo *srcInfo, static Operation *getAncestorInBlock(Operation *op, Block *block) { for (Operation *cur = op; cur; cur = cur->getParentOp()) { - if (cur->getBlock() == block) + if (cur->getBlock() == block) { return cur; + } } return nullptr; } @@ -238,13 +262,16 @@ static bool hasUseAfterOp(Value value, Operation *currentOp) { Block *block = currentOp->getBlock(); for (OpOperand &use : value.getUses()) { Operation *owner = use.getOwner(); - if (owner == currentOp) + if (owner == currentOp) { continue; + } Operation *ancestor = getAncestorInBlock(owner, block); - if (!ancestor) + if (!ancestor) { return true; - if (ancestor != currentOp && currentOp->isBeforeInBlock(ancestor)) + } + if (ancestor != currentOp && currentOp->isBeforeInBlock(ancestor)) { return true; + } } return false; } @@ -254,13 +281,15 @@ static bool hasLaterUseOfSameAddressRange( const Buffer2MemInfoMap &buffer2MemInfoMap) { for (const auto &entry : buffer2MemInfoMap) { // Later reads of the source itself do not make this no-op TMOV a bridge. - if (entry.first == op.getSrc()) + if (entry.first == op.getSrc()) { continue; + } bool sameRange = llvm::any_of(entry.second, [&](const auto &info) { return hasSameConcreteAddressRange(info.get(), dstInfo); }); - if (sameRange && hasUseAfterOp(entry.first, op)) + if (sameRange && hasUseAfterOp(entry.first, op)) { return true; + } } return false; } @@ -271,11 +300,13 @@ static bool hasPlainTMovSemantics(TMovOp op) { } static bool hasCompatibleIdentityTypes(TMovOp op) { - if (op.getSrc().getType() != op.getDst().getType()) + if (op.getSrc().getType() != op.getDst().getType()) { return false; + } for (OpResult result : op->getResults()) { - if (result.getType() != op.getDst().getType()) + if (result.getType() != op.getDst().getType()) { return false; + } } return true; } @@ -305,13 +336,15 @@ isIdentityTMovByMemInfo(TMovOp op, const Buffer2MemInfoMap &buffer2MemInfoMap) { Value src = op.getSrc(); Value dst = op.getDst(); - if (!isStaticallyAddressableValue(src) || !isStaticallyAddressableValue(dst)) + if (!isStaticallyAddressableValue(src) || !isStaticallyAddressableValue(dst)) { return false; + } const BaseMemInfo *srcInfo = getSingleMemInfo(buffer2MemInfoMap, src); const BaseMemInfo *dstInfo = getSingleMemInfo(buffer2MemInfoMap, dst); - if (!hasSameConcreteAddressRange(srcInfo, dstInfo)) + if (!hasSameConcreteAddressRange(srcInfo, dstInfo)) { return false; + } return isDeadDstTMov(op) && !hasLaterUseOfSameAddressRange(op, dstInfo, buffer2MemInfoMap); @@ -327,8 +360,9 @@ struct PTORemoveIdentityTMovPass func.walk([&](TMovOp op) { if (!hasPlainTMovSemantics(op) || !hasCompatibleIdentityTypes(op) || - touchesLowPrecisionElement(op)) + touchesLowPrecisionElement(op)) { return; + } if (op.getSrc() == op.getDst()) { identityMoves.push_back(op); return; @@ -345,14 +379,16 @@ struct PTORemoveIdentityTMovPass translator.Build(); for (TMovOp op : memInfoCandidates) { - if (isIdentityTMovByMemInfo(op, buffer2MemInfoMap)) + if (isIdentityTMovByMemInfo(op, buffer2MemInfoMap)) { identityMoves.push_back(op); + } } } for (TMovOp op : identityMoves) { - for (OpResult result : op->getResults()) + for (OpResult result : op->getResults()) { result.replaceAllUsesWith(op.getDst()); + } op.erase(); } } diff --git a/lib/PTO/Transforms/PTORemoveRedundantBarrier.cpp b/lib/PTO/Transforms/PTORemoveRedundantBarrier.cpp index 4ab5ae5110..40d24c9e6d 100644 --- a/lib/PTO/Transforms/PTORemoveRedundantBarrier.cpp +++ b/lib/PTO/Transforms/PTORemoveRedundantBarrier.cpp @@ -8,19 +8,19 @@ #include "PTO/IR/PTO.h" #include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Pass/Pass.h" #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" #include - + using namespace mlir; using namespace mlir::pto; - + namespace { - + // ========================================================== // 更严格的活跃性分析 // ========================================================== @@ -29,12 +29,15 @@ namespace { // Wait 和 Set 不算作实质性操作。 // 只有真正消耗计算或带宽的指令才算"活跃"。 bool isResourceOp(Operation *op, Attribute targetPipe) { - if (auto loadOp = dyn_cast(op)) + if (auto loadOp = dyn_cast(op)) { return pto::PipeAttr::get(op->getContext(), pto::PIPE::PIPE_MTE2) == targetPipe; - if (auto storeOp = dyn_cast(op)) + } + if (auto storeOp = dyn_cast(op)) { return pto::PipeAttr::get(op->getContext(), pto::PIPE::PIPE_MTE3) == targetPipe; - if (auto addfOp = dyn_cast(op)) + } + if (auto addfOp = dyn_cast(op)) { return pto::PipeAttr::get(op->getContext(), pto::PIPE::PIPE_V) == targetPipe; + } return false; } @@ -44,11 +47,15 @@ bool isPipeUsedInRegion(Region ®ion, Attribute targetPipe) { for (Block &block : region) { for (Operation &op : block) { // 1. 如果是实质性操作，返回 True - if (isResourceOp(&op, targetPipe)) return true; - + if (isResourceOp(&op, targetPipe)) { + return true; + } + // 2. 递归检查嵌套 (if/for) for (Region &nestedRegion : op.getRegions()) { - if (isPipeUsedInRegion(nestedRegion, targetPipe)) return true; + if (isPipeUsedInRegion(nestedRegion, targetPipe)) { + return true; + } } } } @@ -60,15 +67,19 @@ bool isPipeUsedInRegion(Region ®ion, Attribute targetPipe) { // 如果一个 Pipe 后面只剩 Wait，说明它已经完成了工作，发给它的信号是多余的。 static bool hasPipelineActivityAfterOp(Operation *parentOp, Attribute targetPipe) { Block *parentBlock = parentOp ? parentOp->getBlock() : nullptr; - if (!parentBlock) + if (!parentBlock) { return false; + } for (auto it = std::next(parentOp->getIterator()); it != parentBlock->end(); ++it) { - if (isResourceOp(&*it, targetPipe)) + if (isResourceOp(&*it, targetPipe)) { return true; - if (it->getNumRegions() > 0) + } + if (it->getNumRegions() > 0) { return true; - if (isa(&*it)) + } + if (isa(&*it)) { return false; + } } return false; } @@ -76,34 +87,40 @@ static bool hasPipelineActivityAfterOp(Operation *parentOp, Attribute targetPipe bool isPipelineActiveFuture(Block *block, Block::iterator startIt, Attribute targetPipe) { for (auto it = startIt; it != block->end(); ++it) { Operation *op = &*it; - - // 1. 遇到实质性操作 -> 活跃 - if (isResourceOp(op, targetPipe)) return true; - + +// 1. 遇到实质性操作 -> 活跃 + if (isResourceOp(op, targetPipe)) { + return true; + } + // [注意] 这里故意跳过了 WaitOp 的检查。 // WaitOp 只是同步原语，不代表该 Pipeline 在"干活"。 - + // 2. 递归检查嵌套区域 (scf.if, scf.for) for (Region ®ion : op->getRegions()) { - if (isPipeUsedInRegion(region, targetPipe)) return true; + if (isPipeUsedInRegion(region, targetPipe)) { + return true; + } } - + // 3. 处理 Terminator (跨 Block 检查) if (op->hasTrait()) { // 如果是 Return，肯定死了 - if (isa(op)) return false; + if (isa(op)) { + return false; + } return hasPipelineActivityAfterOp(block->getParentOp(), targetPipe); } } return false; } - + // ========================================================== // Pass 实现 // ========================================================== struct PTORemoveRedundantBarrierPass : public PassWrapper> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PTORemoveRedundantBarrierPass) - + void runOnOperation() override { func::FuncOp func = getOperation(); MLIRContext *ctx = &getContext(); @@ -111,11 +128,17 @@ struct PTORemoveRedundantBarrierPass : public PassWrapper Attribute { - if (isa(op)) return attrMTE2; - if (isa(op)) return attrMTE3; - if (isa(op)) return attrVec; + if (isa(op)) { + return attrMTE2; + } + if (isa(op)) { + return attrMTE3; + } + if (isa(op)) { + return attrVec; + } return {}; }; @@ -130,8 +153,9 @@ struct PTORemoveRedundantBarrierPass : public PassWrappergetAttrOfType("src_pipe"); auto dstAttr = op->getAttrOfType("dst_pipe"); - if (!srcAttr || !dstAttr) + if (!srcAttr || !dstAttr) { return false; + } src = srcAttr; dst = dstAttr; return true; @@ -147,31 +171,32 @@ struct PTORemoveRedundantBarrierPass : public PassWrappergetName().getStringRef(); if (opName == "pto.wait_flag_dyn" || opName == "pto.wait_flag_d") { auto dstAttr = op->getAttrOfType("dst_pipe"); - if (!dstAttr) + if (!dstAttr) { return false; + } dst = dstAttr; return true; } return false; }; - + llvm::SmallVector opsToErase; - + func.walk([&](Block *block) { // 记录 Block 内脏状态 (Intra-Block Dirty State) // 用于判断是否需要发广播 llvm::DenseSet intraPipeDirtySet; - + for (auto it = block->begin(); it != block->end(); ++it) { Operation *op = &*it; Attribute pipe = getOpPipe(op); - + // === 1. 状态更新 === if (pipe) { intraPipeDirtySet.insert(pipe); continue; } - + // === 2. Barrier 消除 === if (auto barrierOp = dyn_cast(op)) { Attribute bPipe = barrierOp.getPipe(); @@ -203,7 +228,7 @@ struct PTORemoveRedundantBarrierPass : public PassWrappererase(); + for (Operation *op : opsToErase) { + op->erase(); + } } }; diff --git a/lib/PTO/Transforms/PTOResolveBufferSelect.cpp b/lib/PTO/Transforms/PTOResolveBufferSelect.cpp index 44c49a4a2c..02947709ca 100644 --- a/lib/PTO/Transforms/PTOResolveBufferSelect.cpp +++ b/lib/PTO/Transforms/PTOResolveBufferSelect.cpp @@ -42,20 +42,25 @@ using namespace mlir; namespace { static uint64_t alignUp(uint64_t value, uint64_t align) { - if (align == 0) + if (align == 0) { return value; + } return ((value + align - 1) / align) * align; } static Value ensureI64(Value value, IRRewriter &rewriter, Location loc) { - if (!value) + if (!value) { return {}; - if (value.getType().isInteger(64)) + } + if (value.getType().isInteger(64)) { return value; - if (value.getType().isIndex()) + } + if (value.getType().isIndex()) { return rewriter.create(loc, rewriter.getI64Type(), value); - if (isa(value.getType())) + } + if (isa(value.getType())) { return rewriter.create(loc, rewriter.getI64Type(), value); + } return {}; } @@ -78,8 +83,9 @@ static bool getTilePointerStrides(pto::TileBufType type, int64_t &rowStride, } unsigned elemBytes = pto::getPTOStorageElemByteSize(type.getElementType()); - if (elemBytes == 0) + if (elemBytes == 0) { return false; + } int64_t innerRows = 1; int64_t innerCols = 1; int32_t fractal = config.getSFractalSize().getInt(); @@ -100,8 +106,9 @@ static bool getTilePointerStrides(pto::TileBufType type, int64_t &rowStride, } if (bl == 1) { - if (sl != 1) + if (sl != 1) { return false; + } rowStride = innerCols; colStride = shape[0] + @@ -167,8 +174,9 @@ static Value computeTileAddress(Value value, IRRewriter &rewriter, Value elements = rewriter.create(loc, row, col); int64_t elemBytes = static_cast( pto::getPTOStorageElemByteSize(sourceType.getElementType())); - if (elemBytes == 0) + if (elemBytes == 0) { return {}; + } Value byteScale = rewriter.create(loc, elemBytes, 64); Value bytes = rewriter.create(loc, elements, byteScale); return rewriter.create(loc, base, bytes); @@ -231,8 +239,9 @@ static LogicalResult resolveTileNativeSubviews(ModuleOp module, // A tile function argument is a symbolic runtime-bound handle. Keep its // subview tile-native; only planned local roots can be normalized to an // addressed alloc_tile here. - if (!addr) + if (!addr) { continue; + } pto::TileBufType physicalType = getSubviewPhysicalType(op); auto alloc = rewriter.create( op.getLoc(), physicalType, addr, @@ -246,12 +255,14 @@ static LogicalResult resolveTileNativeSubviews(ModuleOp module, static FailureOr getStaticSlotBytes(pto::TileBufType slotType) { uint64_t elemBytes = pto::getPTOStorageElemByteSize(slotType.getElementType()); - if (elemBytes == 0) + if (elemBytes == 0) { return failure(); + } uint64_t bytes = elemBytes; for (int64_t dim : slotType.getShape()) { - if (dim == ShapedType::kDynamic) + if (dim == ShapedType::kDynamic) { return failure(); + } bytes *= static_cast(dim); } return bytes; @@ -263,8 +274,9 @@ static LogicalResult getMultiTileAddresses(pto::AllocMultiTileOp alloc, uint32_t count = alloc.getResult().getType().getCount(); if (auto planned = alloc->getAttrOfType( pto::kPtoMultiBufferAddrsAttrName)) { - if (planned.size() != count) + if (planned.size() != count) { return alloc.emitError("planned address count does not match slot count"); + } for (int64_t address : planned.asArrayRef()) addrs.push_back(rewriter.create( alloc.getLoc(), address, 64)); @@ -272,13 +284,15 @@ static LogicalResult getMultiTileAddresses(pto::AllocMultiTileOp alloc, } Value base = alloc.getAddr(); - if (!base) + if (!base) { return alloc.emitError( "has neither a level3 base address nor planner-assigned slot addresses"); + } auto slotBytes = getStaticSlotBytes(alloc.getResult().getType().getSlotType()); - if (failed(slotBytes)) + if (failed(slotBytes)) { return alloc.emitError( "requires a static slot shape and known element byte size"); + } uint64_t slotStride = alignUp(*slotBytes, @@ -309,8 +323,9 @@ static LogicalResult resolveTileNativeMultiGets(ModuleOp module, IRRewriter rewriter(ctx); rewriter.setInsertionPoint(op); SmallVector addrs; - if (failed(getMultiTileAddresses(alloc, rewriter, addrs))) + if (failed(getMultiTileAddresses(alloc, rewriter, addrs))) { return failure(); + } Value selectedAddr; IntegerAttr constSlotAttr; @@ -340,9 +355,10 @@ static LogicalResult resolveTileNativeMultiGets(ModuleOp module, SmallVector allocs; module.walk([&](pto::AllocMultiTileOp op) { allocs.push_back(op); }); for (pto::AllocMultiTileOp alloc : allocs) { - if (!alloc.getResult().use_empty()) + if (!alloc.getResult().use_empty()) { return alloc.emitError( "has unsupported uses after resolving pto.multi_tile_get"); + } alloc.erase(); } return success(); diff --git a/lib/PTO/Transforms/PTOResolveReservedBuffersPass.cpp b/lib/PTO/Transforms/PTOResolveReservedBuffersPass.cpp index 2cc999a8e2..3231432741 100644 --- a/lib/PTO/Transforms/PTOResolveReservedBuffersPass.cpp +++ b/lib/PTO/Transforms/PTOResolveReservedBuffersPass.cpp @@ -130,8 +130,9 @@ static std::optional getPipePeerKey(Value localAddr, auto peerFunc = lookupPeerFuncAcrossContainer(importOp.getOperation(), importOp.getPeerFuncAttr()); - if (!peerFunc) + if (!peerFunc) { return std::nullopt; + } return PipePeerKey{getFuncSymbol(peerFunc), importOp.getName().str(), 0}; } @@ -176,11 +177,13 @@ static LogicalResult collectPeerAwareInit(InitOpT initOp, } auto recordAddr = [&](Value addr, int8_t effectiveDirMask) { - if (!addr) + if (!addr) { return false; + } auto key = getPipePeerKey(addr, info.funcOp); - if (!key) + if (!key) { return false; + } key->dirMask = effectiveDirMask; keyedInits[*key].push_back(info.op); return true; @@ -195,10 +198,12 @@ static LogicalResult collectPeerAwareInit(InitOpT initOp, recorded = recordAddr(getLocalAddrOperand(initOp), info.dirMask); } - if (recorded) + if (recorded) { initInfos.push_back(info); - if (recorded || getFlagBaseAttr(initOp)) + } + if (recorded || getFlagBaseAttr(initOp)) { return success(); + } return initOp.emitOpError( "requires local_addr to come from pto.reserve_buffer or " @@ -249,8 +254,9 @@ buildPeerAwareComponents(const SmallVectorImpl &initInfos, for (const auto &it : keyedInits) { SmallVector uniqueOps; for (Operation *op : it.second) { - if (std::find(uniqueOps.begin(), uniqueOps.end(), op) == uniqueOps.end()) + if (std::find(uniqueOps.begin(), uniqueOps.end(), op) == uniqueOps.end()) { uniqueOps.push_back(op); + } } for (size_t i = 0; i < uniqueOps.size(); ++i) { for (size_t j = i + 1; j < uniqueOps.size(); ++j) { @@ -263,8 +269,9 @@ buildPeerAwareComponents(const SmallVectorImpl &initInfos, SmallVector components; llvm::SmallPtrSet visited; for (const PipeInitInfo &rootInfo : initInfos) { - if (!visited.insert(rootInfo.op).second) + if (!visited.insert(rootInfo.op).second) { continue; + } SmallVector stack{rootInfo.op}; PipeComponent component; @@ -272,8 +279,9 @@ buildPeerAwareComponents(const SmallVectorImpl &initInfos, Operation *current = stack.pop_back_val(); component.ops.push_back(current); for (Operation *neighbor : adjacency[current]) { - if (visited.insert(neighbor).second) + if (visited.insert(neighbor).second) { stack.push_back(neighbor); + } } } @@ -327,10 +335,11 @@ buildPeerAwareComponents(const SmallVectorImpl &initInfos, // producer/consumer functions. Keep the smallest observed id only as a // stable component sort key; cross-function pairing is determined by // the peer buffer contract instead of frontend id equality. - if (component.frontendId) + if (component.frontendId) { component.frontendId = std::min(*component.frontendId, *frontendId); - else + } else { component.frontendId = *frontendId; + } } } @@ -408,8 +417,9 @@ static FailureOr chooseFlagBaseForComponent(const PipeComponent &compon nextCandidate = std::max(nextCandidate, alignToEven(used.end)); } } - if (!conflict) + if (!conflict) { break; + } candidateBase = nextCandidate; } @@ -451,8 +461,9 @@ struct PTOResolveReservedBuffersPass if (failed(chosenBaseOr)) return failure(); auto flagBaseAttr = builder.getI32IntegerAttr(*chosenBaseOr); - for (Operation *op : component.ops) + for (Operation *op : component.ops) { setFlagBaseAttr(op, flagBaseAttr); + } } return success(); @@ -523,8 +534,9 @@ struct PTOResolveReservedBuffersPass } } - for (Operation *op : eraseOps) + for (Operation *op : eraseOps) { op->erase(); + } return success(); } diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 4b246c494b..1a96b468f1 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -249,10 +249,11 @@ buildDefaultLastUseTileSlotOrder(Operation *op) { for (OpOperand &operand : op->getOpOperands()) { if (!isa(operand.get().getType())) continue; - if (isDpsInitOperand(operand)) + if (isDpsInitOperand(operand)) { dpsInitTileOperands.push_back(operand.getOperandNumber()); - else + } else { nonDpsTileOperands.push_back(operand.getOperandNumber()); + } } // Most tile intrinsics lower as `CALLEE(dst, src0, src1, ...)`. When an op @@ -421,39 +422,53 @@ static bool isF8E8M0ElemType(Type elemTy) { } static std::string getEmitCScalarTypeToken(Type elemTy) { - if (pto::isPTOFloat8E4M3LikeType(elemTy)) + if (pto::isPTOFloat8E4M3LikeType(elemTy)) { return "float8_e4m3_t"; - if (pto::isPTOFloat8E5M2LikeType(elemTy)) + } + if (pto::isPTOFloat8E5M2LikeType(elemTy)) { return "float8_e5m2_t"; - if (isF8E8M0ElemType(elemTy)) + } + if (isF8E8M0ElemType(elemTy)) { return "float8_e8m0_t"; - if (isa(elemTy)) + } + if (isa(elemTy)) { return "hifloat8_t"; - if (isa(elemTy)) + } + if (isa(elemTy)) { return "float4_e1m2x2_t"; - if (isa(elemTy)) + } + if (isa(elemTy)) { return "float4_e2m1x2_t"; - if (elemTy.isF16()) + } + if (elemTy.isF16()) { return "half"; - if (elemTy.isBF16()) + } + if (elemTy.isBF16()) { return "bfloat16_t"; - if (elemTy.isF32()) + } + if (elemTy.isF32()) { return "float"; - if (elemTy.isF64()) + } + if (elemTy.isF64()) { return "double"; - if (elemTy.isInteger(8)) + } + if (elemTy.isInteger(8)) { return (elemTy.isSignlessInteger(8) || elemTy.isSignedInteger(8)) ? "int8_t" : "uint8_t"; - if (elemTy.isInteger(16)) + } + if (elemTy.isInteger(16)) { return (elemTy.isSignlessInteger(16) || elemTy.isSignedInteger(16)) ? "int16_t" : "uint16_t"; - if (elemTy.isInteger(32)) + } + if (elemTy.isInteger(32)) { return (elemTy.isSignlessInteger(32) || elemTy.isSignedInteger(32)) ? "int32_t" : "uint32_t"; - if (elemTy.isInteger(64)) + } + if (elemTy.isInteger(64)) { return cast(elemTy).isUnsigned() ? "uint64_t" : "int64_t"; + } return "float"; } @@ -835,14 +850,24 @@ class PTOToEmitCTypeConverter : public TypeConverter { // 1. 基本类型 (f32, i32, index) // --------------------------------------------------------- addConversion([Ctx](FloatType type) -> Type { - if (pto::isPTOFloat8E4M3LikeType(type)) + if (pto::isPTOFloat8E4M3LikeType(type)) { return emitc::OpaqueType::get(Ctx, "float8_e4m3_t"); - if (pto::isPTOFloat8E5M2LikeType(type)) + } + if (pto::isPTOFloat8E5M2LikeType(type)) { return emitc::OpaqueType::get(Ctx, "float8_e5m2_t"); - if (type.isF32()) return emitc::OpaqueType::get(Ctx, "float"); - if (type.isF16()) return emitc::OpaqueType::get(Ctx, "half"); - if (type.isBF16()) return emitc::OpaqueType::get(Ctx, "bfloat16_t"); - if (type.isF64()) return emitc::OpaqueType::get(Ctx, "double"); + } + if (type.isF32()) { + return emitc::OpaqueType::get(Ctx, "float"); + } + if (type.isF16()) { + return emitc::OpaqueType::get(Ctx, "half"); + } + if (type.isBF16()) { + return emitc::OpaqueType::get(Ctx, "bfloat16_t"); + } + if (type.isF64()) { + return emitc::OpaqueType::get(Ctx, "double"); + } llvm::errs() << "[Debug] Unsupported FloatType: " << type << "\n"; return Type{}; }); @@ -1042,15 +1067,21 @@ class PTOToEmitCTypeConverter : public TypeConverter { // --------------------------------------------------------- addConversion([this](FunctionType type) -> Type { SmallVector inputs; - if (failed(convertTypes(type.getInputs(), inputs))) return Type{}; + if (failed(convertTypes(type.getInputs(), inputs))) { + return Type{}; + } SmallVector results; - if (failed(convertTypes(type.getResults(), results))) return Type{}; + if (failed(convertTypes(type.getResults(), results))) { + return Type{}; + } return FunctionType::get(type.getContext(), inputs, results); }); auto materializeCast = [](OpBuilder &Builder, Type ResultType, ValueRange Inputs, Location Loc) -> Value { - if (Inputs.size() != 1) return Value(); + if (Inputs.size() != 1) { + return Value(); + } return Builder.create(Loc, ResultType, Inputs[0]).getResult(0); }; @@ -1526,10 +1557,11 @@ adaptCallOperandForEmitC(const TypeConverter *typeConverter, } else if (auto memrefTy = dyn_cast(originalCalleeArgTy)) { elemTy = memrefTy.getElementType(); if (auto asAttr = - dyn_cast_or_null(memrefTy.getMemorySpace())) + dyn_cast_or_null(memrefTy.getMemorySpace())) { as = asAttr.getAddressSpace(); - else + } else { as = pto::AddressSpace::GM; + } } if (elemTy && as) { @@ -2619,10 +2651,10 @@ struct ArithMulExtendedToEmitC : public OpConversionPattern { Type lowDstTy = newResultTypes[0]; Type highDstTy = newResultTypes[1]; - Type wideTy = isUnsigned ? (Type)getWiderUnsignedIntOpaqueType(rewriter.getContext(), - bitWidth) - : (Type)getWiderSignedIntOpaqueType(rewriter.getContext(), - bitWidth); + Type wideTy = isUnsigned ? static_cast(getWiderUnsignedIntOpaqueType(rewriter.getContext(), + bitWidth)) + : static_cast(getWiderSignedIntOpaqueType(rewriter.getContext(), + bitWidth)); Value lhsWide; Value rhsWide; @@ -3539,12 +3571,22 @@ enum class KernelKind { VecAdd, Matmul, Unknown }; bool hasAdd = false; bool hasMM = false; f.walk([&](Operation *op) { - if (isa(op)) hasAdd = true; - if (isa(op)) hasMM = true; - if (isa(op)) hasMM = true; + if (isa(op)) { + hasAdd = true; + } + if (isa(op)) { + hasMM = true; + } + if (isa(op)) { + hasMM = true; + } }); - if (hasMM) return KernelKind::Matmul; - if (hasAdd) return KernelKind::VecAdd; + if (hasMM) { + return KernelKind::Matmul; + } + if (hasAdd) { + return KernelKind::VecAdd; + } return KernelKind::Unknown; } @@ -3556,12 +3598,14 @@ enum class KernelKind { VecAdd, Matmul, Unknown }; auto readShape2D = [&](memref::SubViewOp sv, int &d0, int &d1) { auto resTy = mlir::cast(sv.getResult().getType()); if (resTy.getRank() == 2 && resTy.hasStaticShape()) { - d0 = (int)resTy.getDimSize(0); - d1 = (int)resTy.getDimSize(1); + d0 = static_cast(resTy.getDimSize(0)); + d1 = static_cast(resTy.getDimSize(1)); } }; - if (subs.empty()) return; + if (subs.empty()) { + return; + } int a0=32, a1=32; readShape2D(subs[0], a0, a1); @@ -3893,7 +3937,7 @@ struct SubviewToEmitCPattern : public OpConversionPattern { // B. 获取 Stride (用于指针计算) Value strideVal = mkIndex(1); - if (i < (int)sourceStrides.size()) { + if (i < static_cast(sourceStrides.size())) { strideVal = ofrToEmitCValue(sourceStrides[i]); } @@ -4002,11 +4046,12 @@ struct SubviewToEmitCPattern : public OpConversionPattern { shapeParamsVec.push_back(resShape[i]); } // size 值：优先从 op.getMixedSizes() 取（可动态/静态），否则退化为类型里的静态 shape。 - if (i < (int)mixedSizes.size()) + if (i < static_cast(mixedSizes.size())) { sizeValues.push_back(ofrToEmitCValue(mixedSizes[i])); - else + } else { sizeValues.push_back( mkIndex(resShape[i] == ShapedType::kDynamic ? 1 : resShape[i])); + } } // 3. 生成 Stride 模板参数 + 运行时 stride 值（考虑 subview step） @@ -4017,9 +4062,9 @@ struct SubviewToEmitCPattern : public OpConversionPattern { auto subViewSteps = op.getMixedStrides(); for (int i = 0; i < rank; ++i) { OpFoldResult srcStrideOfr = - (i < (int)sourceStrides.size()) ? sourceStrides[i] - : rewriter.getIndexAttr(1); - OpFoldResult stepOfr = (i < (int)subViewSteps.size()) + (i < static_cast(sourceStrides.size())) ? sourceStrides[i] + : rewriter.getIndexAttr(1); + OpFoldResult stepOfr = (i < static_cast(subViewSteps.size())) ? subViewSteps[i] : rewriter.getIndexAttr(1); @@ -4036,13 +4081,14 @@ struct SubviewToEmitCPattern : public OpConversionPattern { Value srcV = ofrToEmitCValue(srcStrideOfr); Value stepV = ofrToEmitCValue(stepOfr); // 尽量避免乘以 1 生成冗余指令 - if (stepStatic && *stepStatic == 1) + if (stepStatic && *stepStatic == 1) { strideValues.push_back(srcV); - else if (srcStatic && *srcStatic == 1) + } else if (srcStatic && *srcStatic == 1) { strideValues.push_back(stepV); - else + } else { strideValues.push_back( rewriter.create(loc, indexTy, srcV, stepV)); + } } // 3.1 右对齐到 5 维：shape 补 1；已有维度继承原 stride； @@ -4124,13 +4170,15 @@ struct SubviewToEmitCPattern : public OpConversionPattern { layoutTag = 2; // NZ } else { bool isRow = finalStride[4] == 1; - for (int i = 3; i >= 0; --i) + for (int i = 3; i >= 0; --i) { isRow &= (finalStride[i] == multiplyOrDynamic(finalStride[i + 1], finalShape[i + 1])); + } bool isCol = finalStride[0] == 1; - for (int i = 0; i < 4; ++i) + for (int i = 0; i < 4; ++i) { isCol &= (finalStride[i + 1] == multiplyOrDynamic(finalStride[i], finalShape[i])); + } if (isCol) layoutTag = 1; // DN else @@ -4514,15 +4562,16 @@ static Value materializeGlobalTensorDataPointer( return value; Type elemType; - if (auto tvTy = dyn_cast(sourceType)) + if (auto tvTy = dyn_cast(sourceType)) { elemType = tvTy.getElementType(); - else if (auto partitionTy = - dyn_cast(sourceType)) + } else if (auto partitionTy = + dyn_cast(sourceType)) { elemType = partitionTy.getElementType(); - else if (auto memrefTy = dyn_cast(sourceType)) + } else if (auto memrefTy = dyn_cast(sourceType)) { elemType = memrefTy.getElementType(); - else + } else { return value; + } auto *ctx = rewriter.getContext(); std::string elemTypeStr = getElemTypeStringForGT(elemType); @@ -5460,13 +5509,15 @@ struct PTOFenceToEmitC : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { (void)adaptor; if (op.getScope().getScope() != pto::FenceScope::GM && - op.getScope().getScope() != pto::FenceScope::All) + op.getScope().getScope() != pto::FenceScope::All) { return rewriter.notifyMatchFailure(op, "unsupported fence scope"); + } - if (isInVectorKernel(op)) + if (isInVectorKernel(op)) { emitPipeBarrier(rewriter, op.getLoc(), "PIPE_ALL"); - else + } else { emitConservativeGmFencePipeDrains(rewriter, op.getLoc()); + } emitDsbDdr(rewriter, op.getLoc()); rewriter.eraseOp(op); return success(); @@ -6139,7 +6190,6 @@ struct PTOGetBlockIdxToEmitC LogicalResult matchAndRewrite(mlir::pto::GetBlockIdxOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - rewriter.replaceOpWithNewOp( op, op.getType(), "get_block_idx", ValueRange{}, ArrayAttr{}, ArrayAttr{}); @@ -6156,7 +6206,6 @@ struct PTOGetBlockNumToEmitC LogicalResult matchAndRewrite(mlir::pto::GetBlockNumOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - rewriter.replaceOpWithNewOp( op, op.getType(), "get_block_num", ValueRange{}, ArrayAttr{}, ArrayAttr{}); @@ -6173,7 +6222,6 @@ struct PTOGetSubBlockIdxToEmitC LogicalResult matchAndRewrite(mlir::pto::GetSubBlockIdxOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - rewriter.replaceOpWithNewOp( op, op.getType(), "get_subblockid", ValueRange{}, ArrayAttr{}, ArrayAttr{}); @@ -6190,7 +6238,6 @@ struct PTOGetSubBlockNumToEmitC LogicalResult matchAndRewrite(mlir::pto::GetSubBlockNumOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - rewriter.replaceOpWithNewOp( op, op.getType(), "get_subblockdim", ValueRange{}, ArrayAttr{}, ArrayAttr{}); @@ -6853,20 +6900,22 @@ struct PTOInitializeL2G2LPipeToEmitC Value c2vBuf = zero; Value v2cBuf = zero; - if (op.getDirMask() == 1) + if (op.getDirMask() == 1) { c2vBuf = localAddr ? localAddr : zero; - else if (op.getDirMask() == 2) + } else if (op.getDirMask() == 2) { v2cBuf = localAddr ? localAddr : zero; - else if (op.getDirMask() == 3) { + } else if (op.getDirMask() == 3) { if (localAddr) { - if (!op.getPeerLocalAddr()) + if (!op.getPeerLocalAddr()) { return rewriter.notifyMatchFailure( op, "bidirectional l2g2l pipe requires peer local buffer"); + } c2vBuf = localAddr; v2cBuf = peelUnrealized(adaptor.getPeerLocalAddr()); } - } else + } else { return rewriter.notifyMatchFailure(op, "unsupported dir_mask"); + } rewriter.replaceOpWithNewOp( op, TypeRange{emitPipeTy}, *tpipeTok, ArrayAttr{}, ArrayAttr{}, @@ -6905,15 +6954,16 @@ struct PTOInitializeL2LPipeToEmitC Value c2vBuf = zero; Value v2cBuf = zero; - if (op.getDirMask() == 1) + if (op.getDirMask() == 1) { c2vBuf = localAddr; - else if (op.getDirMask() == 2) + } else if (op.getDirMask() == 2) { v2cBuf = localAddr; - else if (op.getDirMask() == 3) { + } else if (op.getDirMask() == 3) { c2vBuf = localAddr; v2cBuf = peelUnrealized(adaptor.getPeerLocalAddr()); - } else + } else { return rewriter.notifyMatchFailure(op, "unsupported dir_mask"); + } rewriter.replaceOpWithNewOp( op, TypeRange{emitPipeTy}, *tpipeTok, ArrayAttr{}, ArrayAttr{}, @@ -8076,12 +8126,13 @@ struct PTOGetTensorViewMetadataToEmitC : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Type sourceType = op.getTensorView().getType(); int64_t rank = 0; - if (auto type = dyn_cast(sourceType)) + if (auto type = dyn_cast(sourceType)) { rank = type.getRank(); - else if (auto type = dyn_cast(sourceType)) + } else if (auto type = dyn_cast(sourceType)) { rank = type.getRank(); - else + } else { return rewriter.notifyMatchFailure(op, "expected PTO tensor view"); + } Value result = getRuntimeGlobalTensorMetadata( rewriter, op.getLoc(), peelUnrealized(adaptor.getTensorView()), @@ -8098,7 +8149,7 @@ static LogicalResult getStaticTensorViewStrides( strides.clear(); if (auto makeView = source.getDefiningOp()) { - if ((int64_t)makeView.getStrides().size() != rank) + if (static_cast(makeView.getStrides().size()) != rank) return failure(); for (Value strideValue : makeView.getStrides()) { auto cst = getStaticIndexLikeValue(strideValue); @@ -8115,7 +8166,7 @@ static LogicalResult getStaticTensorViewStrides( StringRef token = opaqueTy.getValue(); if ((parseIntegerTemplateList(token, "pto::Stride<", stride5D) || parseIntegerTemplateList(token, "Stride<", stride5D)) && - (int64_t)stride5D.size() >= rank) { + static_cast(stride5D.size()) >= rank) { strides.append(stride5D.end() - rank, stride5D.end()); return success(); } @@ -8625,10 +8676,11 @@ struct ReinterpretCastToEmitC : public OpConversionPattern { std::string scalarTok = "int32_t"; if (auto it = dyn_cast(op->getOperand(0).getType())) { bool isUnsigned = it.isUnsigned(); - if (it.getWidth() == 16) + if (it.getWidth() == 16) { scalarTok = isUnsigned ? "uint16_t" : "int16_t"; - else + } else { scalarTok = isUnsigned ? "uint32_t" : "int32_t"; + } } // descending -> "0"/"1" @@ -9864,7 +9917,6 @@ struct PTOFillPadToEmitC : public OpConversionPattern { //===----------------------------------------------------------------------===// [[maybe_unused]] static std::string maskPatternTok(mlir::pto::MaskPatternAttr a) { - auto v = a.getValue(); // enum return (std::string("pto::MaskPattern::") + mlir::pto::stringifyMaskPattern(v).str()); } @@ -11199,10 +11251,11 @@ struct PTORowExpandAddToEmitC : public OpConversionPattern Value dst = peelUnrealized(adaptor.getDst()); SmallVector operands; - if (tmp) + if (tmp) { operands.assign({dst, src0, src1, tmp}); - else + } else { operands.assign({dst, src0, src1}); + } rewriter.create( loc, TypeRange{}, "TROWEXPANDADD", /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, @@ -11227,10 +11280,11 @@ struct PTORowExpandExpdifToEmitC Value tmp = op.getTmp() ? peelUnrealized(adaptor.getTmp()) : Value(); SmallVector operands; - if (tmp) + if (tmp) { operands.assign({dst, src0, src1, tmp}); - else + } else { operands.assign({dst, src0, src1}); + } rewriter.create( loc, TypeRange{}, "TROWEXPANDEXPDIF", /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, @@ -11247,10 +11301,11 @@ static void replaceOrEraseWithOpaqueCallAndReturnDst(Operation *op, Value dst, ArrayAttr templateArgs, ConversionPatternRewriter &rewriter) { createLastUseAwareOpaqueCall(rewriter, op, TypeRange{}, callee, args, ArrayAttr{}, templateArgs); - if (op->getNumResults() == 1) + if (op->getNumResults() == 1) { rewriter.replaceOp(op, dst); - else + } else { rewriter.eraseOp(op); + } } // ---------- TOp ---------- @@ -11431,10 +11486,11 @@ struct PTORowExpandDivToEmitC : public OpConversionPattern Value tmp = op.getTmp() ? peelUnrealized(adaptor.getTmp()) : Value(); SmallVector operands; - if (tmp) + if (tmp) { operands.assign({dst, src0, src1, tmp}); - else + } else { operands.assign({dst, src0, src1}); + } ArrayAttr templateArgs; if (op.getPrecisionType() != pto::DivPrecision::Default) { StringRef precisionTok; @@ -11471,10 +11527,11 @@ struct PTORowExpandMulToEmitC : public OpConversionPattern Value tmp = op.getTmp() ? peelUnrealized(adaptor.getTmp()) : Value(); SmallVector operands; - if (tmp) + if (tmp) { operands.assign({dst, src0, src1, tmp}); - else + } else { operands.assign({dst, src0, src1}); + } createLastUseAwareOpaqueCall(rewriter, op.getOperation(), TypeRange{}, "TROWEXPANDMUL", operands); @@ -11500,10 +11557,11 @@ struct PTORowExpandSubToEmitC : public OpConversionPattern Value tmp = op.getTmp() ? peelUnrealized(adaptor.getTmp()) : Value(); SmallVector operands; - if (tmp) + if (tmp) { operands.assign({dst, src0, src1, tmp}); - else + } else { operands.assign({dst, src0, src1}); + } rewriter.create( loc, TypeRange{}, "TROWEXPANDSUB", /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, @@ -11527,10 +11585,11 @@ struct PTORowExpandMaxToEmitC : public OpConversionPattern Value tmp = op.getTmp() ? peelUnrealized(adaptor.getTmp()) : Value(); SmallVector operands; - if (tmp) + if (tmp) { operands.assign({dst, src0, src1, tmp}); - else + } else { operands.assign({dst, src0, src1}); + } rewriter.create( loc, TypeRange{}, "TROWEXPANDMAX", /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, @@ -11554,10 +11613,11 @@ struct PTORowExpandMinToEmitC : public OpConversionPattern Value tmp = op.getTmp() ? peelUnrealized(adaptor.getTmp()) : Value(); SmallVector operands; - if (tmp) + if (tmp) { operands.assign({dst, src0, src1, tmp}); - else + } else { operands.assign({dst, src0, src1}); + } rewriter.create( loc, TypeRange{}, "TROWEXPANDMIN", /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, @@ -11952,10 +12012,11 @@ struct PTOSORT32SToEmitC : public OpConversionPattern { Value tmp = op.getTmp() ? peelUnrealized(adaptor.getTmp()) : Value(); SmallVector operands; - if (tmp) + if (tmp) { operands.assign({dst, src, idx, tmp}); - else + } else { operands.assign({dst, src, idx}); + } rewriter.create( loc, TypeRange{}, "TSORT32", /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, @@ -12255,14 +12316,15 @@ struct PTOPrintOpToEmitC : public OpConversionPattern { fmt = "%f"; std::string quoted = "\""; for (char c : fmt) { - if (c == '"' || c == '\\') + if (c == '"' || c == '\\') { quoted += '\\'; - else if (c == '\n') + } else if (c == '\n') { quoted += "\\n"; - else if (c == '\t') + } else if (c == '\t') { quoted += "\\t"; - else + } else { quoted += c; + } } quoted += "\""; diff --git a/lib/PTO/Transforms/PTOVPTOPtrBoundary.cpp b/lib/PTO/Transforms/PTOVPTOPtrBoundary.cpp index 1ba2d136a7..65269d8dfc 100644 --- a/lib/PTO/Transforms/PTOVPTOPtrBoundary.cpp +++ b/lib/PTO/Transforms/PTOVPTOPtrBoundary.cpp @@ -30,12 +30,14 @@ namespace { static Type convertVPTOBoundaryMemRefType(Type type) { auto memrefType = dyn_cast(type); - if (!memrefType) + if (!memrefType) { return type; + } auto memorySpace = dyn_cast_or_null(memrefType.getMemorySpace()); - if (!memorySpace) + if (!memorySpace) { return {}; + } return pto::PtrType::get(type.getContext(), memrefType.getElementType(), memorySpace); } @@ -56,29 +58,33 @@ static LogicalResult eraseDeadVPTOMemRefScaffold(ModuleOp module) { trivialCasts.push_back(castOp); return; } - if (castOp->use_empty()) + if (castOp->use_empty()) { deadOps.push_back(op); + } return; } - if (!op->use_empty()) + if (!op->use_empty()) { return; + } if (isa(op)) deadOps.push_back(op); }); for (pto::CastPtrOp castOp : trivialCasts) { - if (!castOp->getBlock()) + if (!castOp->getBlock()) { continue; + } castOp.getResult().replaceAllUsesWith(castOp.getInput()); castOp.erase(); erasedAny = true; } for (Operation *op : deadOps) { - if (!op->getBlock()) + if (!op->getBlock()) { continue; + } op->erase(); erasedAny = true; } @@ -88,23 +94,29 @@ static LogicalResult eraseDeadVPTOMemRefScaffold(ModuleOp module) { static Type getVPTOBufferElementType(Value value) { Type type = value.getType(); - if (auto tileType = dyn_cast(type)) + if (auto tileType = dyn_cast(type)) { return tileType.getElementType(); - if (auto memrefType = dyn_cast(type)) + } + if (auto memrefType = dyn_cast(type)) { return memrefType.getElementType(); - if (auto ptrType = dyn_cast(type)) + } + if (auto ptrType = dyn_cast(type)) { return ptrType.getElementType(); + } return {}; } static Attribute getVPTOBufferMemorySpace(Value value) { Type type = value.getType(); - if (auto tileType = dyn_cast(type)) + if (auto tileType = dyn_cast(type)) { return tileType.getMemorySpace(); - if (auto memrefType = dyn_cast(type)) + } + if (auto memrefType = dyn_cast(type)) { return memrefType.getMemorySpace(); - if (auto ptrType = dyn_cast(type)) + } + if (auto ptrType = dyn_cast(type)) { return ptrType.getMemorySpace(); + } return {}; } @@ -125,21 +137,25 @@ static LogicalResult canonicalizeBoundaryCastPtrOps(ModuleOp module, llvm::raw_ostream *diagOS) { SmallVector castsToRewrite; module.walk([&](pto::CastPtrOp castOp) { - if (!isa(castOp.getInput().getType())) + if (!isa(castOp.getInput().getType())) { return; - if (!isa(castOp.getResult().getType())) + } + if (!isa(castOp.getResult().getType())) { return; + } castsToRewrite.push_back(castOp); }); PatternRewriter rewriter(module.getContext()); for (pto::CastPtrOp castOp : castsToRewrite) { - if (!castOp->getBlock()) + if (!castOp->getBlock()) { continue; + } auto resultType = dyn_cast(castOp.getResult().getType()); - if (!resultType) + if (!resultType) { continue; + } rewriter.setInsertionPoint(castOp); Value ptrValue = pto::materializeBufferPointer( @@ -166,8 +182,9 @@ static LogicalResult canonicalizeSupportedVPTOBufferLikeOps( ModuleOp module, llvm::raw_ostream *diagOS) { SmallVector opsToRewrite; module.walk([&](Operation *op) { - if (isSupportedVPTOBufferLikeBoundaryOp(op)) + if (isSupportedVPTOBufferLikeBoundaryOp(op)) { opsToRewrite.push_back(op); + } }); PatternRewriter rewriter(module.getContext()); @@ -213,8 +230,9 @@ static LogicalResult canonicalizeSupportedVPTOBufferLikeOps( newOperands.push_back(ptrValue); } - if (!changed) + if (!changed) { continue; + } OperationState state(op->getLoc(), op->getName().getStringRef()); state.addOperands(newOperands); @@ -235,8 +253,9 @@ struct PTOVPTOPtrBoundaryPass void runOnOperation() override { ModuleOp module = getOperation(); - if (failed(pto::convertVPTOEmissionBoundaryToPtr(module, &llvm::errs()))) + if (failed(pto::convertVPTOEmissionBoundaryToPtr(module, &llvm::errs()))) { signalPassFailure(); + } } }; @@ -248,13 +267,15 @@ LogicalResult mlir::pto::convertVPTOEmissionBoundaryToPtr( // function ABI keeps only the same-space base pointer, while shape/stride // state remains in SSA. Body-level op canonicalization is added on top of // this entry rewrite in follow-up tasks. - if (failed(eraseDeadVPTOMemRefScaffold(module))) + if (failed(eraseDeadVPTOMemRefScaffold(module))) { return failure(); + } bool sawFailure = false; for (func::FuncOp func : module.getOps()) { - if (func.isExternal()) + if (func.isExternal()) { continue; + } FunctionType functionType = func.getFunctionType(); SmallVector newInputs(functionType.getInputs().begin(), @@ -263,8 +284,9 @@ LogicalResult mlir::pto::convertVPTOEmissionBoundaryToPtr( for (auto [idx, inputType] : llvm::enumerate(functionType.getInputs())) { auto memrefType = dyn_cast(inputType); - if (!memrefType) + if (!memrefType) { continue; + } Type newType = convertVPTOBoundaryMemRefType(inputType); if (!newType) { @@ -284,8 +306,9 @@ LogicalResult mlir::pto::convertVPTOEmissionBoundaryToPtr( for (Operation *user : users) { if (auto cast = dyn_cast(user)) { - if (cast.getInput() != arg) + if (cast.getInput() != arg) { continue; + } if (cast.getResult().getType() == newType) { cast.getResult().replaceAllUsesWith(arg); cast.erase(); @@ -300,8 +323,9 @@ LogicalResult mlir::pto::convertVPTOEmissionBoundaryToPtr( continue; } - if (isSupportedVPTOBufferLikeBoundaryOp(user)) + if (isSupportedVPTOBufferLikeBoundaryOp(user)) { continue; + } if (diagOS) { *diagOS << "VPTO emission-boundary ptr rewrite failed: argument " @@ -315,8 +339,9 @@ LogicalResult mlir::pto::convertVPTOEmissionBoundaryToPtr( } for (Type resultType : functionType.getResults()) { - if (!isa(resultType)) + if (!isa(resultType)) { continue; + } if (diagOS) *diagOS << "VPTO emission-boundary ptr rewrite failed: memref result " "is unsupported for " @@ -330,14 +355,17 @@ LogicalResult mlir::pto::convertVPTOEmissionBoundaryToPtr( } } - if (sawFailure) + if (sawFailure) { return failure(); + } - if (failed(canonicalizeBoundaryCastPtrOps(module, diagOS))) + if (failed(canonicalizeBoundaryCastPtrOps(module, diagOS))) { return failure(); + } - if (failed(canonicalizeSupportedVPTOBufferLikeOps(module, diagOS))) + if (failed(canonicalizeSupportedVPTOBufferLikeOps(module, diagOS))) { return failure(); + } return eraseDeadVPTOMemRefScaffold(module); } diff --git a/lib/PTO/Transforms/PTOValidateIntToPtrUses.cpp b/lib/PTO/Transforms/PTOValidateIntToPtrUses.cpp index 990b81d4ee..eaec47c1d9 100644 --- a/lib/PTO/Transforms/PTOValidateIntToPtrUses.cpp +++ b/lib/PTO/Transforms/PTOValidateIntToPtrUses.cpp @@ -27,8 +27,9 @@ using namespace mlir::pto; static bool isAllowedIntToPtrUse(Value ptr, OpOperand &use) { Operation *user = use.getOwner(); - if (isa(user)) + if (isa(user)) { return use.getOperandNumber() == 0 && user->getOperand(0) == ptr; + } return false; } diff --git a/lib/PTO/Transforms/PTOValidatePhysicalSectionBoundaries.cpp b/lib/PTO/Transforms/PTOValidatePhysicalSectionBoundaries.cpp index 3d0801bdc7..1036b3f7bc 100644 --- a/lib/PTO/Transforms/PTOValidatePhysicalSectionBoundaries.cpp +++ b/lib/PTO/Transforms/PTOValidatePhysicalSectionBoundaries.cpp @@ -32,41 +32,48 @@ namespace { static Operation *getNearestPhysicalSection(Operation *op) { for (Operation *current = op; current; current = current->getParentOp()) { - if (isa(current)) + if (isa(current)) { return current; + } } return nullptr; } static Operation *getNearestPhysicalSection(BlockArgument argument) { - if (!argument) + if (!argument) { return nullptr; + } return getNearestPhysicalSection(argument.getOwner()->getParentOp()); } static Operation *getDefiningPhysicalSection(Value value) { - if (auto blockArgument = dyn_cast(value)) + if (auto blockArgument = dyn_cast(value)) { return getNearestPhysicalSection(blockArgument); - if (Operation *definingOp = value.getDefiningOp()) + } + if (Operation *definingOp = value.getDefiningOp()) { return getNearestPhysicalSection(definingOp); + } return nullptr; } static StringRef getPhysicalSectionKind(Operation *section) { - if (!section) + if (!section) { return "outside a physical section"; + } return isa(section) ? "pto.section.cube" : "pto.section.vector"; } static LogicalResult verifyOperandBoundary(Operation *user, OpOperand &operand) { Operation *definingSection = getDefiningPhysicalSection(operand.get()); - if (!definingSection) + if (!definingSection) { return success(); + } Operation *usingSection = getNearestPhysicalSection(user); - if (definingSection == usingSection) + if (definingSection == usingSection) { return success(); + } return user->emitOpError() << "value defined in " << getPhysicalSectionKind(definingSection) @@ -81,8 +88,9 @@ static LogicalResult verifyOperandBoundary(Operation *user, OpOperand &operand) static LogicalResult verifyFunction(func::FuncOp function) { LogicalResult result = success(); function.walk([&](Operation *op) { - if (failed(result)) + if (failed(result)) { return WalkResult::interrupt(); + } for (OpOperand &operand : op->getOpOperands()) { if (failed(verifyOperandBoundary(op, operand))) { result = failure(); @@ -101,13 +109,15 @@ struct PTOValidatePhysicalSectionBoundariesPass ModuleOp module = getOperation(); LogicalResult result = success(); module.walk([&](func::FuncOp function) { - if (failed(result)) + if (failed(result)) { return WalkResult::interrupt(); + } result = verifyFunction(function); return failed(result) ? WalkResult::interrupt() : WalkResult::advance(); }); - if (failed(result)) + if (failed(result)) { signalPassFailure(); + } } }; diff --git a/lib/PTO/Transforms/PTOValidateVMIIR.cpp b/lib/PTO/Transforms/PTOValidateVMIIR.cpp index 84aafb1399..9f9cd6ff67 100644 --- a/lib/PTO/Transforms/PTOValidateVMIIR.cpp +++ b/lib/PTO/Transforms/PTOValidateVMIIR.cpp @@ -41,8 +41,9 @@ namespace { bool isVMIType(Type type) { return isa(type); } bool containsVMIType(Type type) { - if (isVMIType(type)) + if (isVMIType(type)) { return true; + } if (auto functionType = dyn_cast(type)) { return llvm::any_of(functionType.getInputs(), @@ -52,51 +53,63 @@ bool containsVMIType(Type type) { }); } - if (auto shapedType = dyn_cast(type)) + if (auto shapedType = dyn_cast(type)) { return containsVMIType(shapedType.getElementType()); + } return false; } bool containsVMIType(Attribute attr) { - if (!attr) + if (!attr) { return false; + } - if (auto typeAttr = dyn_cast(attr)) - if (containsVMIType(typeAttr.getValue())) + if (auto typeAttr = dyn_cast(attr)) { + if (containsVMIType(typeAttr.getValue())) { return true; + } + } - if (auto typedAttr = dyn_cast(attr)) - if (containsVMIType(typedAttr.getType())) + if (auto typedAttr = dyn_cast(attr)) { + if (containsVMIType(typedAttr.getType())) { return true; + } + } - if (auto arrayAttr = dyn_cast(attr)) + if (auto arrayAttr = dyn_cast(attr)) { return llvm::any_of(arrayAttr, [](Attribute element) { return containsVMIType(element); }); + } - if (auto dictAttr = dyn_cast(attr)) + if (auto dictAttr = dyn_cast(attr)) { return llvm::any_of(dictAttr, [](NamedAttribute namedAttr) { return containsVMIType(namedAttr.getValue()); }); + } return false; } bool isSurfaceVMIType(Type type) { - if (auto vregType = dyn_cast(type)) + if (auto vregType = dyn_cast(type)) { return !vregType.getLayout(); - if (auto maskType = dyn_cast(type)) + } + if (auto maskType = dyn_cast(type)) { return maskType.isPred() && !maskType.getLayout(); + } return false; } bool isLayoutAssignedVMIType(Type type) { - if (auto vregType = dyn_cast(type)) + if (auto vregType = dyn_cast(type)) { return static_cast(vregType.getLayoutAttr()); - if (auto maskType = dyn_cast(type)) + } + if (auto maskType = dyn_cast(type)) { return maskType.getLayoutAttr() && VMIMaskType::isConcreteGranularity(maskType.getGranularity()); + } return false; } @@ -132,16 +145,18 @@ bool hasVMIType(Operation *op) { return true; for (Region ®ion : op->getRegions()) { for (Block &block : region) { - if (llvm::any_of(block.getArgumentTypes(), isVMIType)) + if (llvm::any_of(block.getArgumentTypes(), isVMIType)) { return true; + } } } return false; } void mirrorDiagnostic(llvm::raw_ostream *diagOS, Twine message) { - if (diagOS) + if (diagOS) { *diagOS << message << "\n"; + } } LogicalResult emitInvariant(Operation *op, llvm::raw_ostream *diagOS, @@ -171,8 +186,9 @@ LogicalResult emitLayoutSupportContract(Operation *op, bool printedAny = false; auto printValueType = [&](StringRef kind, int64_t index, Type type) { - if (!isVMIType(type)) + if (!isVMIType(type)) { return; + } if (!printedAny) { os << "; VMI types:"; printedAny = true; @@ -180,10 +196,12 @@ LogicalResult emitLayoutSupportContract(Operation *op, os << " " << kind << "#" << index << "=" << type; }; - for (auto [index, operand] : llvm::enumerate(op->getOperands())) + for (auto [index, operand] : llvm::enumerate(op->getOperands())) { printValueType("operand", static_cast(index), operand.getType()); - for (auto [index, result] : llvm::enumerate(op->getResults())) + } + for (auto [index, result] : llvm::enumerate(op->getResults())) { printValueType("result", static_cast(index), result.getType()); + } os.flush(); return emitLayoutContract(op, diagOS, text); @@ -200,8 +218,9 @@ emitHelperMaterializationContract(Operation *helper, Type sourceType, " has no registered materialization support: " + reason); }; - if (helper->getNumResults() != 1 || !helper->getResult(0).hasOneUse()) + if (helper->getNumResults() != 1 || !helper->getResult(0).hasOneUse()) { return emitFallback(); + } OpOperand &use = *helper->getResult(0).use_begin(); Operation *requester = use.getOwner(); @@ -234,20 +253,26 @@ LogicalResult verifyBoundaryType(Operation *owner, Type type, LogicalResult verifyBoundaryTypeTree(Operation *owner, Type type, llvm::raw_ostream *diagOS) { - if (failed(verifyBoundaryType(owner, type, diagOS))) + if (failed(verifyBoundaryType(owner, type, diagOS))) { return failure(); + } if (auto functionType = dyn_cast(type)) { - for (Type input : functionType.getInputs()) - if (failed(verifyBoundaryTypeTree(owner, input, diagOS))) + for (Type input : functionType.getInputs()) { + if (failed(verifyBoundaryTypeTree(owner, input, diagOS))) { return failure(); - for (Type result : functionType.getResults()) - if (failed(verifyBoundaryTypeTree(owner, result, diagOS))) + } + } + for (Type result : functionType.getResults()) { + if (failed(verifyBoundaryTypeTree(owner, result, diagOS))) { return failure(); + } + } } - if (auto shapedType = dyn_cast(type)) + if (auto shapedType = dyn_cast(type)) { return verifyBoundaryTypeTree(owner, shapedType.getElementType(), diagOS); + } return success(); } @@ -265,21 +290,27 @@ LogicalResult verifyLayoutAssignedType(Operation *owner, Type type, LogicalResult verifyLayoutAssignedTypeTree(Operation *owner, Type type, llvm::raw_ostream *diagOS) { - if (failed(verifyLayoutAssignedType(owner, type, diagOS))) + if (failed(verifyLayoutAssignedType(owner, type, diagOS))) { return failure(); + } if (auto functionType = dyn_cast(type)) { - for (Type input : functionType.getInputs()) - if (failed(verifyLayoutAssignedTypeTree(owner, input, diagOS))) + for (Type input : functionType.getInputs()) { + if (failed(verifyLayoutAssignedTypeTree(owner, input, diagOS))) { return failure(); - for (Type result : functionType.getResults()) - if (failed(verifyLayoutAssignedTypeTree(owner, result, diagOS))) + } + } + for (Type result : functionType.getResults()) { + if (failed(verifyLayoutAssignedTypeTree(owner, result, diagOS))) { return failure(); + } + } } - if (auto shapedType = dyn_cast(type)) + if (auto shapedType = dyn_cast(type)) { return verifyLayoutAssignedTypeTree(owner, shapedType.getElementType(), diagOS); + } return success(); } @@ -288,28 +319,37 @@ template LogicalResult verifyAttributeTypes(Operation *owner, Attribute attr, llvm::raw_ostream *diagOS, TypeVerifier verifyType) { - if (!attr) + if (!attr) { return success(); + } - if (auto typeAttr = dyn_cast(attr)) - if (failed(verifyType(owner, typeAttr.getValue(), diagOS))) + if (auto typeAttr = dyn_cast(attr)) { + if (failed(verifyType(owner, typeAttr.getValue(), diagOS))) { return failure(); + } + } - if (auto typedAttr = dyn_cast(attr)) - if (failed(verifyType(owner, typedAttr.getType(), diagOS))) + if (auto typedAttr = dyn_cast(attr)) { + if (failed(verifyType(owner, typedAttr.getType(), diagOS))) { return failure(); + } + } if (auto arrayAttr = dyn_cast(attr)) { - for (Attribute element : arrayAttr) - if (failed(verifyAttributeTypes(owner, element, diagOS, verifyType))) + for (Attribute element : arrayAttr) { + if (failed(verifyAttributeTypes(owner, element, diagOS, verifyType))) { return failure(); + } + } } if (auto dictAttr = dyn_cast(attr)) { - for (NamedAttribute namedAttr : dictAttr) + for (NamedAttribute namedAttr : dictAttr) { if (failed(verifyAttributeTypes(owner, namedAttr.getValue(), diagOS, - verifyType))) + verifyType))) { return failure(); + } + } } return success(); @@ -321,45 +361,58 @@ bool isFunctionTypeAttr(Operation *op, NamedAttribute attr) { LogicalResult verifyNoHiddenVMIAttributeType(Operation *op, NamedAttribute attr, llvm::raw_ostream *diagOS) { - if (isFunctionTypeAttr(op, attr)) + if (isFunctionTypeAttr(op, attr)) { return success(); - if (containsVMIType(attr.getValue())) + } + if (containsVMIType(attr.getValue())) { return emitInvariant(op, diagOS, "VMI type appears in a non-signature attribute"); + } return success(); } LogicalResult verifyOperationTypes(Operation *op, llvm::raw_ostream *diagOS) { if (auto funcOp = dyn_cast(op)) { FunctionType functionType = funcOp.getFunctionType(); - for (Type type : functionType.getInputs()) - if (failed(verifyBoundaryTypeTree(op, type, diagOS))) + for (Type type : functionType.getInputs()) { + if (failed(verifyBoundaryTypeTree(op, type, diagOS))) { return failure(); - for (Type type : functionType.getResults()) - if (failed(verifyBoundaryTypeTree(op, type, diagOS))) + } + } + for (Type type : functionType.getResults()) { + if (failed(verifyBoundaryTypeTree(op, type, diagOS))) { return failure(); + } + } } - for (Type type : op->getOperandTypes()) - if (failed(verifyBoundaryTypeTree(op, type, diagOS))) + for (Type type : op->getOperandTypes()) { + if (failed(verifyBoundaryTypeTree(op, type, diagOS))) { return failure(); - for (Type type : op->getResultTypes()) - if (failed(verifyBoundaryTypeTree(op, type, diagOS))) + } + } + for (Type type : op->getResultTypes()) { + if (failed(verifyBoundaryTypeTree(op, type, diagOS))) { return failure(); + } + } for (Region ®ion : op->getRegions()) { for (Block &block : region) { for (Type type : block.getArgumentTypes()) { - if (failed(verifyBoundaryTypeTree(op, type, diagOS))) + if (failed(verifyBoundaryTypeTree(op, type, diagOS))) { return failure(); + } } } } for (NamedAttribute attr : op->getAttrs()) { - if (failed(verifyNoHiddenVMIAttributeType(op, attr, diagOS))) + if (failed(verifyNoHiddenVMIAttributeType(op, attr, diagOS))) { return failure(); + } if (failed(verifyAttributeTypes(op, attr.getValue(), diagOS, - verifyBoundaryTypeTree))) + verifyBoundaryTypeTree))) { return failure(); + } } return success(); } @@ -368,34 +421,45 @@ LogicalResult verifyLayoutAssignedOperationTypes(Operation *op, llvm::raw_ostream *diagOS) { if (auto funcOp = dyn_cast(op)) { FunctionType functionType = funcOp.getFunctionType(); - for (Type type : functionType.getInputs()) - if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) + for (Type type : functionType.getInputs()) { + if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) { return failure(); - for (Type type : functionType.getResults()) - if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) + } + } + for (Type type : functionType.getResults()) { + if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) { return failure(); + } + } } - for (Type type : op->getOperandTypes()) - if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) + for (Type type : op->getOperandTypes()) { + if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) { return failure(); - for (Type type : op->getResultTypes()) - if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) + } + } + for (Type type : op->getResultTypes()) { + if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) { return failure(); + } + } for (Region ®ion : op->getRegions()) { for (Block &block : region) { for (Type type : block.getArgumentTypes()) { - if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) + if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) { return failure(); + } } } } for (NamedAttribute attr : op->getAttrs()) { - if (failed(verifyNoHiddenVMIAttributeType(op, attr, diagOS))) + if (failed(verifyNoHiddenVMIAttributeType(op, attr, diagOS))) { return failure(); + } if (failed(verifyAttributeTypes(op, attr.getValue(), diagOS, - verifyLayoutAssignedTypeTree))) + verifyLayoutAssignedTypeTree))) { return failure(); + } } return success(); } @@ -408,19 +472,23 @@ LogicalResult verifyLayoutSemanticSupport(Operation *op, LogicalResult verifyOperationBoundary(Operation *op, llvm::raw_ostream *diagOS) { - if (failed(verifyOperationTypes(op, diagOS))) + if (failed(verifyOperationTypes(op, diagOS))) { return failure(); + } - if (!hasVMIType(op)) + if (!hasVMIType(op)) { return success(); + } - if (isVMIHelperOp(op)) + if (isVMIHelperOp(op)) { return emitInvariant( op, diagOS, "VMI helper op appears before layout assignment or VMI-to-VPTO"); + } - if (isVMISemanticOp(op) || isStructuralOp(op)) + if (isVMISemanticOp(op) || isStructuralOp(op)) { return success(); + } return emitInvariant(op, diagOS, "VMI typed value is used by a non-VMI semantic op"); @@ -429,25 +497,30 @@ LogicalResult verifyOperationBoundary(Operation *op, LogicalResult verifyLayoutAssignedOperation(Operation *op, llvm::raw_ostream *diagOS, bool verifyHelperSupports = true) { - if (failed(verifyLayoutAssignedOperationTypes(op, diagOS))) + if (failed(verifyLayoutAssignedOperationTypes(op, diagOS))) { return failure(); + } - if (!hasVMIType(op)) + if (!hasVMIType(op)) { return success(); + } if (isVMIHelperOp(op)) { - if (isVMILayoutHelperOp(op)) + if (isVMILayoutHelperOp(op)) { return verifyHelperSupports ? verifyLayoutHelperSupport(op, diagOS) : success(); + } return emitInvariant( op, diagOS, "VMI pack/unpack helper appears before VMI-to-VPTO physicalization"); } - if (isVMISemanticOp(op)) + if (isVMISemanticOp(op)) { return verifyLayoutSemanticSupport(op, diagOS); - if (isStructuralOp(op)) + } + if (isStructuralOp(op)) { return success(); + } return emitInvariant(op, diagOS, "VMI typed value is used by a non-VMI semantic op"); @@ -461,9 +534,10 @@ LogicalResult verifyLayoutHelperSupport(Operation *op, auto sourceType = cast(ensure.getSource().getType()); auto resultType = cast(ensure.getResult().getType()); std::string reason; - if (failed(supports.getEnsureLayoutFact(sourceType, resultType, &reason))) + if (failed(supports.getEnsureLayoutFact(sourceType, resultType, &reason))) { return emitHelperMaterializationContract( op, sourceType, resultType, "pto.vmi.ensure_layout", reason, diagOS); + } return success(); } @@ -472,10 +546,11 @@ LogicalResult verifyLayoutHelperSupport(Operation *op, auto resultType = cast(ensure.getResult().getType()); std::string reason; if (failed( - supports.getEnsureMaskLayoutFact(sourceType, resultType, &reason))) + supports.getEnsureMaskLayoutFact(sourceType, resultType, &reason))) { return emitHelperMaterializationContract(op, sourceType, resultType, "pto.vmi.ensure_mask_layout", reason, diagOS); + } return success(); } @@ -489,8 +564,9 @@ LogicalResult verifyLayoutSemanticSupport(Operation *op, if (auto store = dyn_cast(op)) { auto valueType = cast(store.getValue().getType()); VMILayoutAttr layout = valueType.getLayoutAttr(); - if (!layout || layout.isContiguous()) + if (!layout || layout.isContiguous()) { return success(); + } std::string reason; if (failed(supports.getStoreLayoutFact(valueType, &reason))) @@ -504,8 +580,9 @@ LogicalResult verifyLayoutSemanticSupport(Operation *op, if (auto load = dyn_cast(op)) { auto resultType = cast(load.getResult().getType()); VMILayoutAttr layout = resultType.getLayoutAttr(); - if (!layout) + if (!layout) { return success(); + } std::string reason; if (failed(supports.getGroupLoadLayoutFact(load, &reason))) @@ -539,8 +616,9 @@ LogicalResult verifyLayoutSemanticSupport(Operation *op, if (auto store = dyn_cast(op)) { auto valueType = cast(store.getValue().getType()); VMILayoutAttr layout = valueType.getLayoutAttr(); - if (!layout || !layout.isGroupSlots()) + if (!layout || !layout.isGroupSlots()) { return success(); + } std::string reason; if (failed(supports.getGroupStoreLayoutFact( @@ -555,8 +633,9 @@ LogicalResult verifyLayoutSemanticSupport(Operation *op, if (auto reduce = dyn_cast(op)) { auto resultType = cast(reduce.getResult().getType()); VMILayoutAttr layout = resultType.getLayoutAttr(); - if (!layout || !layout.isGroupSlots()) + if (!layout || !layout.isGroupSlots()) { return success(); + } std::string reason; if (failed(supports.getGroupReduceAddFSupport(reduce, &reason))) @@ -571,8 +650,9 @@ LogicalResult verifyLayoutSemanticSupport(Operation *op, if (auto reduce = dyn_cast(op)) { auto resultType = cast(reduce.getResult().getType()); VMILayoutAttr layout = resultType.getLayoutAttr(); - if (!layout || !layout.isGroupSlots()) + if (!layout || !layout.isGroupSlots()) { return success(); + } std::string reason; if (failed(supports.getGroupReduceMaxFSupport(reduce, &reason))) @@ -587,8 +667,9 @@ LogicalResult verifyLayoutSemanticSupport(Operation *op, if (auto reduce = dyn_cast(op)) { auto resultType = cast(reduce.getResult().getType()); VMILayoutAttr layout = resultType.getLayoutAttr(); - if (!layout || !layout.isGroupSlots()) + if (!layout || !layout.isGroupSlots()) { return success(); + } std::string reason; if (failed(supports.getGroupReduceMinFSupport(reduce, &reason))) @@ -603,8 +684,9 @@ LogicalResult verifyLayoutSemanticSupport(Operation *op, if (auto reduce = dyn_cast(op)) { auto resultType = cast(reduce.getResult().getType()); VMILayoutAttr layout = resultType.getLayoutAttr(); - if (!layout || !layout.isGroupSlots()) + if (!layout || !layout.isGroupSlots()) { return success(); + } std::string reason; if (failed(supports.getGroupReduceAddISupport(reduce, &reason))) @@ -619,8 +701,9 @@ LogicalResult verifyLayoutSemanticSupport(Operation *op, if (auto reduce = dyn_cast(op)) { auto resultType = cast(reduce.getResult().getType()); VMILayoutAttr layout = resultType.getLayoutAttr(); - if (!layout || !layout.isGroupSlots()) + if (!layout || !layout.isGroupSlots()) { return success(); + } std::string reason; if (failed(supports.getGroupReduceMaxISupport(reduce, &reason))) @@ -635,8 +718,9 @@ LogicalResult verifyLayoutSemanticSupport(Operation *op, if (auto reduce = dyn_cast(op)) { auto resultType = cast(reduce.getResult().getType()); VMILayoutAttr layout = resultType.getLayoutAttr(); - if (!layout || !layout.isGroupSlots()) + if (!layout || !layout.isGroupSlots()) { return success(); + } std::string reason; if (failed(supports.getGroupReduceMinISupport(reduce, &reason))) @@ -651,8 +735,9 @@ LogicalResult verifyLayoutSemanticSupport(Operation *op, if (auto broadcast = dyn_cast(op)) { auto sourceType = cast(broadcast.getSource().getType()); VMILayoutAttr layout = sourceType.getLayoutAttr(); - if (!layout || !layout.isGroupSlots() || layout.getSlots() <= 0) + if (!layout || !layout.isGroupSlots() || layout.getSlots() <= 0) { return success(); + } std::string reason; if (failed(supports.getGroupBroadcastSupport(broadcast, &reason))) @@ -714,8 +799,9 @@ struct PTOValidateVMIIRPass MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PTOValidateVMIIRPass) void runOnOperation() override { - if (failed(validateVMIProducerBoundaryIR(getOperation(), &llvm::errs()))) + if (failed(validateVMIProducerBoundaryIR(getOperation(), &llvm::errs()))) { signalPassFailure(); + } } }; @@ -725,8 +811,9 @@ struct PTOValidateVMILayoutIRPass MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PTOValidateVMILayoutIRPass) void runOnOperation() override { - if (failed(validateVMILayoutAssignedIR(getOperation(), &llvm::errs()))) + if (failed(validateVMILayoutAssignedIR(getOperation(), &llvm::errs()))) { signalPassFailure(); + } } }; @@ -736,8 +823,9 @@ LogicalResult mlir::pto::validateVMIProducerBoundaryIR(ModuleOp module, llvm::raw_ostream *diagOS) { WalkResult result = module.walk([&](Operation *op) { - if (failed(verifyOperationBoundary(op, diagOS))) + if (failed(verifyOperationBoundary(op, diagOS))) { return WalkResult::interrupt(); + } return WalkResult::advance(); }); return failure(result.wasInterrupted()); @@ -746,8 +834,9 @@ mlir::pto::validateVMIProducerBoundaryIR(ModuleOp module, LogicalResult mlir::pto::validateVMILayoutAssignedIR( ModuleOp module, llvm::raw_ostream *diagOS, bool verifyHelperSupports) { WalkResult result = module.walk([&](Operation *op) { - if (failed(verifyLayoutAssignedOperation(op, diagOS, verifyHelperSupports))) + if (failed(verifyLayoutAssignedOperation(op, diagOS, verifyHelperSupports))) { return WalkResult::interrupt(); + } return WalkResult::advance(); }); return failure(result.wasInterrupted()); diff --git a/lib/PTO/Transforms/PTOValidateVPTOIR.cpp b/lib/PTO/Transforms/PTOValidateVPTOIR.cpp index 75ca0c1707..df505f3477 100644 --- a/lib/PTO/Transforms/PTOValidateVPTOIR.cpp +++ b/lib/PTO/Transforms/PTOValidateVPTOIR.cpp @@ -51,21 +51,25 @@ LogicalResult validateVPTOEmissionIR(ModuleOp module, namespace detail { static Operation *getFirstNonConstantLikeOp(Block *block) { - if (!block) + if (!block) { return nullptr; + } for (Operation &op : *block) { - if (!op.hasTrait()) + if (!op.hasTrait()) { return &op; + } } return nullptr; } static bool isOpInRange(Operation *op, Operation *first, Operation *last) { for (Operation *cur = first; cur; cur = cur->getNextNode()) { - if (cur == op) + if (cur == op) { return true; - if (cur == last) + } + if (cur == last) { return false; + } } return false; } @@ -74,14 +78,17 @@ static constexpr int64_t kSimtKeepResumeSlotLimit = 123; static std::optional getSimtKeepResumeRegisterCount(Type type) { if (auto intType = dyn_cast(type)) { - if (intType.getWidth() <= 32) + if (intType.getWidth() <= 32) { return 1; - if (intType.getWidth() == 64) + } + if (intType.getWidth() == 64) { return 2; + } return std::nullopt; } - if (type.isF16() || type.isBF16() || type.isF32()) + if (type.isF16() || type.isBF16() || type.isF32()) { return 1; + } return std::nullopt; } @@ -102,8 +109,9 @@ template static LogicalResult verifySimtKeepResumeSlotRange(OpT op) { std::optional registerCount = getSimtKeepResumeRegisterCount(getSimtKeepResumeValueType(op)); - if (!registerCount) + if (!registerCount) { return success(); + } int64_t slot = op.getSlot(); if (slot < 0 || slot >= kSimtKeepResumeSlotLimit) return op.emitOpError() @@ -127,15 +135,18 @@ static bool overlapsEarlierSimtKeepResumeSlotUse(OpT op, SmallVectorImpl &used) { std::optional registerCount = getSimtKeepResumeRegisterCount(getSimtKeepResumeValueType(op)); - if (!registerCount) + if (!registerCount) { return false; + } int64_t slot = op.getSlot(); for (int64_t word = slot; word < slot + *registerCount; ++word) { - if (llvm::is_contained(used, word)) + if (llvm::is_contained(used, word)) { return true; + } } - for (int64_t word = slot; word < slot + *registerCount; ++word) + for (int64_t word = slot; word < slot + *registerCount; ++word) { used.push_back(word); + } return false; } @@ -144,8 +155,9 @@ static LogicalResult verifyUniqueResumeGroupSlots(ResumeOp current, SmallVector slots; for (Operation *cur = first; cur; cur = cur->getNextNode()) { auto resume = dyn_cast(cur); - if (!resume) + if (!resume) { break; + } if (overlapsEarlierSimtKeepResumeSlotUse(resume, slots) && resume.getOperation() == current.getOperation()) return current.emitOpError() @@ -161,15 +173,17 @@ static LogicalResult verifyUniqueKeepGroupSlots(KeepOp current, SmallVector slots; for (Operation *cur = first; cur; cur = cur->getNextNode()) { auto keep = dyn_cast(cur); - if (!keep) + if (!keep) { break; + } if (overlapsEarlierSimtKeepResumeSlotUse(keep, slots) && keep.getOperation() == current.getOperation()) return current.emitOpError() << "duplicates an earlier slot " << keep.getSlot() << " in the SIMT keep epilogue group"; - if (cur == last) + if (cur == last) { break; + } } return success(); } @@ -203,8 +217,9 @@ class VPTOLegalityHelper { SmallVector getFunctions() { SmallVector funcs; - for (func::FuncOp func : module.getOps()) + for (func::FuncOp func : module.getOps()) { funcs.push_back(func); + } return funcs; } @@ -217,8 +232,9 @@ class VPTOLegalityHelper { } static bool requiresVecScope(Operation *op) { - if (!isPTOp(op)) + if (!isPTOp(op)) { return false; + } return llvm::any_of(op->getOperandTypes(), isLegalityTypedValue) || llvm::any_of(op->getResultTypes(), isLegalityTypedValue); @@ -233,34 +249,40 @@ class VPTOLegalityHelper { } static bool isAnyVectorScopeCarrier(Operation *op) { - if (auto loop = dyn_cast_or_null(op)) + if (auto loop = dyn_cast_or_null(op)) { return isAIVectorScopeCarrier(loop); + } return isDedicatedVecScopeCarrier(op); } static Operation *getEnclosingVectorScopeCarrier(Operation *op) { for (Operation *parent = op ? op->getParentOp() : nullptr; parent; parent = parent->getParentOp()) { - if (isAnyVectorScopeCarrier(parent)) + if (isAnyVectorScopeCarrier(parent)) { return parent; + } } return nullptr; } static std::optional getMaskGranularity(Type type) { auto maskType = dyn_cast(type); - if (!maskType) + if (!maskType) { return std::nullopt; + } return getMaskGranularity(maskType); } static std::optional getMaskGranularity(MaskType type) { - if (type.isB8()) + if (type.isB8()) { return VPTOMaskGranularity::B8; - if (type.isB16()) + } + if (type.isB16()) { return VPTOMaskGranularity::B16; - if (type.isB32()) + } + if (type.isB32()) { return VPTOMaskGranularity::B32; + } return std::nullopt; } @@ -278,17 +300,21 @@ class VPTOLegalityHelper { static std::optional inferMaskGranularityFromType(Type type) { - if (auto vregType = dyn_cast(type)) + if (auto vregType = dyn_cast(type)) { type = vregType.getElementType(); + } - if (type.isF32()) + if (type.isF32()) { return VPTOMaskGranularity::B32; - if (type.isF16() || type.isBF16()) + } + if (type.isF16() || type.isBF16()) { return VPTOMaskGranularity::B16; + } auto intType = dyn_cast(type); - if (!intType) + if (!intType) { return std::nullopt; + } switch (intType.getWidth()) { case 8: @@ -305,21 +331,26 @@ class VPTOLegalityHelper { static std::optional inferMaskGranularityFromFamily(Operation *op) { StringRef mnemonic = getPTOpMnemonic(op); - if (mnemonic.empty()) + if (mnemonic.empty()) { return std::nullopt; + } - if (mnemonic.ends_with("_b8")) + if (mnemonic.ends_with("_b8")) { return VPTOMaskGranularity::B8; - if (mnemonic.ends_with("_b16")) + } + if (mnemonic.ends_with("_b16")) { return VPTOMaskGranularity::B16; - if (mnemonic.ends_with("_b32")) + } + if (mnemonic.ends_with("_b32")) { return VPTOMaskGranularity::B32; + } return std::nullopt; } static VPTOBufferAddressFamily classifyBufferAddressFamily(Operation *op) { - if (!op) + if (!op) { return VPTOBufferAddressFamily::None; + } if (isa(op)) @@ -353,8 +384,9 @@ class VPTOLegalityHelper { static SmallVector collectBufferOperands(Operation *op) { SmallVector bufferOperands; for (OpOperand &operand : op->getOpOperands()) { - if (isBufferLikeValue(operand.get().getType())) + if (isBufferLikeValue(operand.get().getType())) { bufferOperands.push_back(&operand); + } } return bufferOperands; } @@ -365,8 +397,9 @@ class VPTOLegalityHelper { } static StringRef getPTOpMnemonic(Operation *op) { - if (!isPTOp(op)) + if (!isPTOp(op)) { return {}; + } StringRef mnemonic = op->getName().getStringRef(); (void)mnemonic.consume_front("pto."); return mnemonic; @@ -393,8 +426,9 @@ class VPTOLegalityValidator { return failure(); } - if (failed(validateAuthoringRules())) + if (failed(validateAuthoringRules())) { return failure(); + } if (stage == VPTOLegalityStage::Emission && failed(validateEmissionRules())) @@ -405,18 +439,22 @@ class VPTOLegalityValidator { private: LogicalResult validateAuthoringRules() { - if (failed(validateAuthoringFunctionSurface())) + if (failed(validateAuthoringFunctionSurface())) { return failure(); - if (failed(validateAuthoringOperationSurface())) + } + if (failed(validateAuthoringOperationSurface())) { return failure(); + } return success(); } LogicalResult validateEmissionRules() { - if (failed(validateEmissionFunctionSurface())) + if (failed(validateEmissionFunctionSurface())) { return failure(); - if (failed(validateEmissionOperationSurface())) + } + if (failed(validateEmissionOperationSurface())) { return failure(); + } return success(); } @@ -435,8 +473,9 @@ class VPTOLegalityValidator { StringRef vectorRole) { auto actual = VPTOLegalityHelper::getMaskGranularity(maskType); auto expected = VPTOLegalityHelper::inferMaskGranularityFromType(vectorType); - if (!actual || !expected || *actual == *expected) + if (!actual || !expected || *actual == *expected) { return success(); + } return op->emitOpError() << maskRole << " " << maskType << " does not match " << vectorRole @@ -447,18 +486,21 @@ class VPTOLegalityValidator { static std::optional inferVstsMaskGranularityOverride(Operation *op) { Value value; - if (auto vsts = dyn_cast(op)) + if (auto vsts = dyn_cast(op)) { value = vsts.getValue(); - else + } else { return std::nullopt; + } auto valueType = dyn_cast(value.getType()); - if (!valueType) + if (!valueType) { return std::nullopt; + } auto distAttr = op->getAttrOfType("dist"); - if (!distAttr) + if (!distAttr) { return std::nullopt; + } StringRef dist = distAttr.getValue(); auto elementType = valueType.getElementType(); @@ -476,38 +518,45 @@ class VPTOLegalityValidator { } if (dist == "PK_B16") { - if (width == 8) + if (width == 8) { return VPTOMaskGranularity::B16; + } return std::nullopt; } if (dist == "PK_B32") { - if (width == 16) + if (width == 16) { return VPTOMaskGranularity::B32; + } return std::nullopt; } if (dist == "PK_B64") { - if (width == 32) + if (width == 32) { return VPTOMaskGranularity::B32; + } return std::nullopt; } if (dist == "PK4_B32") { - if (width == 8) + if (width == 8) { return VPTOMaskGranularity::B32; + } return std::nullopt; } if (dist == "MRG4CHN_B8") { - if (width == 8) + if (width == 8) { return VPTOMaskGranularity::B32; + } return std::nullopt; } if (dist == "MRG2CHN_B8") { - if (width == 8) + if (width == 8) { return VPTOMaskGranularity::B16; + } return std::nullopt; } if (dist == "MRG2CHN_B16") { - if (width == 16) + if (width == 16) { return VPTOMaskGranularity::B32; + } } return std::nullopt; } @@ -518,8 +567,9 @@ class VPTOLegalityValidator { StringRef rhsRole) { auto lhs = VPTOLegalityHelper::getMaskGranularity(lhsType); auto rhs = VPTOLegalityHelper::getMaskGranularity(rhsType); - if (!lhs || !rhs || *lhs == *rhs) + if (!lhs || !rhs || *lhs == *rhs) { return success(); + } return op->emitOpError() << lhsRole << " " << lhsType << " does not match " << rhsRole << " " << rhsType; @@ -589,8 +639,9 @@ class VPTOLegalityValidator { inferVstsMaskGranularityOverride(op.getOperation())) { auto actual = VPTOLegalityHelper::getMaskGranularity(op.getMask().getType()); - if (!actual || *actual == *expected) + if (!actual || *actual == *expected) { return success(); + } return op.emitOpError() << "mask type " << op.getMask().getType() << " does not match value vector type " @@ -607,8 +658,9 @@ class VPTOLegalityValidator { auto emitForStore = [&](auto storeOp) { Operation *store = storeOp.getOperation(); auto distAttr = store->getAttrOfType("dist"); - if (!distAttr) + if (!distAttr) { return; + } StringRef dist = distAttr.getValue(); if (dist == "MRG4CHN_B8" || dist == "MRG2CHN_B8" || dist == "MRG2CHN_B16") @@ -718,8 +770,9 @@ class VPTOLegalityValidator { static LogicalResult validatePredicateMovementContract( PredicateMovementOp op) { auto expected = VPTOLegalityHelper::inferMaskGranularityFromFamily(op); - if (!expected) + if (!expected) { return success(); + } if (failed(validateSameMaskGranularity(op, op.getLhs().getType(), "lhs mask type", @@ -732,12 +785,14 @@ class VPTOLegalityValidator { failed(validateSameMaskGranularity(op, op.getLhs().getType(), "lhs mask type", op.getHigh().getType(), - "high mask type"))) + "high mask type"))) { return failure(); + } auto lhs = VPTOLegalityHelper::getMaskGranularity(op.getLhs().getType()); - if (!lhs || *lhs == *expected) + if (!lhs || *lhs == *expected) { return success(); + } return op.emitOpError() << "predicate movement family requires " @@ -750,8 +805,9 @@ class VPTOLegalityValidator { StringRef resultRole) { auto expected = VPTOLegalityHelper::inferMaskGranularityFromFamily(op); auto actual = VPTOLegalityHelper::getMaskGranularity(resultType); - if (!expected || !actual || *expected == *actual) + if (!expected || !actual || *expected == *actual) { return success(); + } return op->emitOpError() << "family suffix requires " << resultRole << " to be " @@ -777,13 +833,15 @@ class VPTOLegalityValidator { return llvm::TypeSwitch(op) .Case([](VreluOp concreteOp) { auto vecType = dyn_cast(concreteOp.getInput().getType()); - if (!vecType) + if (!vecType) { return success(); + } Type elemType = vecType.getElementType(); if (auto intType = dyn_cast(elemType)) { - if (intType.getWidth() == 32 && !intType.isUnsigned()) + if (intType.getWidth() == 32 && !intType.isUnsigned()) { return success(); + } } else if (elemType.isF16() || elemType.isF32()) { return success(); } @@ -866,8 +924,9 @@ class VPTOLegalityValidator { [&](StringRef attrName, int64_t upperBound, StringRef description) -> LogicalResult { Attribute attr = func->getAttr(attrName); - if (!attr) + if (!attr) { return success(); + } auto intAttr = dyn_cast(attr); if (!intAttr || !intAttr.getType().isSignlessInteger(32)) @@ -896,11 +955,13 @@ class VPTOLegalityValidator { if (failed(validatePositiveI32FuncAttr(pto::kPTOSimtMaxThreadsAttrName, 2048, "SIMT max threads")) || failed(validatePositiveI32FuncAttr(pto::kPTOSimtMaxRegistersAttrName, - 128, "SIMT max registers"))) + 128, "SIMT max registers"))) { return failure(); + } - if (!func->hasAttr(pto::kPTOSimtEntryAttrName)) + if (!func->hasAttr(pto::kPTOSimtEntryAttrName)) { continue; + } WalkResult walkResult = func.walk([&](StoreVfSimtInfoOp op) { op.emitOpError() @@ -910,13 +971,15 @@ class VPTOLegalityValidator { "instead"; return WalkResult::interrupt(); }); - if (walkResult.wasInterrupted()) + if (walkResult.wasInterrupted()) { return failure(); + } } WalkResult keepResumeWalk = helper.getModule().walk([&](Operation *op) { - if (!isa(op)) + if (!isa(op)) { return WalkResult::advance(); + } func::FuncOp func = op->getParentOfType(); if (!func || !func->hasAttr(pto::kPTOSimtEntryAttrName)) { op->emitOpError() @@ -926,8 +989,9 @@ class VPTOLegalityValidator { } Block *block = op->getBlock(); if (auto resume = dyn_cast(op)) { - if (failed(verifySimtKeepResumeSlotRange(resume))) + if (failed(verifySimtKeepResumeSlotRange(resume))) { return WalkResult::interrupt(); + } Operation *first = getFirstNonConstantLikeOp(block); if (!first || !isa(first)) { op->emitOpError() @@ -937,8 +1001,9 @@ class VPTOLegalityValidator { } bool found = false; for (Operation *cur = first; cur; cur = cur->getNextNode()) { - if (!isa(cur)) + if (!isa(cur)) { break; + } if (cur == op) { found = true; break; @@ -950,12 +1015,14 @@ class VPTOLegalityValidator { "constant-like operations"; return WalkResult::interrupt(); } - if (failed(verifyUniqueResumeGroupSlots(resume, first))) + if (failed(verifyUniqueResumeGroupSlots(resume, first))) { return WalkResult::interrupt(); + } } if (auto keep = dyn_cast(op)) { - if (failed(verifySimtKeepResumeSlotRange(keep))) + if (failed(verifySimtKeepResumeSlotRange(keep))) { return WalkResult::interrupt(); + } Operation *terminator = block ? block->getTerminator() : nullptr; if (!terminator || !isa(terminator)) { op->emitOpError() @@ -964,8 +1031,9 @@ class VPTOLegalityValidator { } Operation *cur = terminator->getPrevNode(); - while (cur && isa(cur)) + while (cur && isa(cur)) { cur = cur->getPrevNode(); + } Operation *lastKeep = cur; if (!lastKeep || !isa(lastKeep)) { op->emitOpError() @@ -977,8 +1045,9 @@ class VPTOLegalityValidator { Operation *firstKeep = lastKeep; while (Operation *prev = firstKeep->getPrevNode()) { - if (!isa(prev)) + if (!isa(prev)) { break; + } firstKeep = prev; } if (!isOpInRange(op, firstKeep, lastKeep)) { @@ -994,8 +1063,9 @@ class VPTOLegalityValidator { } return WalkResult::advance(); }); - if (keepResumeWalk.wasInterrupted()) + if (keepResumeWalk.wasInterrupted()) { return failure(); + } return success(); } @@ -1004,27 +1074,32 @@ class VPTOLegalityValidator { helper.getModule().walk([&](arith::ConstantOp constant) { Type resultType = constant.getType(); Type elementType = resultType; - if (auto vectorType = dyn_cast(resultType)) + if (auto vectorType = dyn_cast(resultType)) { elementType = vectorType.getElementType(); - if (!pto::isPTOFloat8Type(elementType)) + } + if (!pto::isPTOFloat8Type(elementType)) { return WalkResult::advance(); + } constant.emitOpError() << "does not support directly constructed FP8 constants in " "the VPTO backend; produce FP8 values with pto.convert"; return WalkResult::interrupt(); }); - if (constantWalkResult.wasInterrupted()) + if (constantWalkResult.wasInterrupted()) { return failure(); + } WalkResult loopWalkResult = helper.getModule().walk([&](scf::ForOp loop) { - if (!VPTOLegalityHelper::isAIVectorScopeCarrier(loop)) + if (!VPTOLegalityHelper::isAIVectorScopeCarrier(loop)) { return WalkResult::advance(); + } Operation *parentScope = VPTOLegalityHelper::getEnclosingVectorScopeCarrier(loop); - if (!parentScope) + if (!parentScope) { return WalkResult::advance(); + } if (isa(parentScope)) { loop.emitOpError() << "does not allow nested scf.for with '" @@ -1037,35 +1112,41 @@ class VPTOLegalityValidator { "pto.vecscope/pto.strict_vecscope"; return WalkResult::interrupt(); }); - if (loopWalkResult.wasInterrupted()) + if (loopWalkResult.wasInterrupted()) { return failure(); + } WalkResult vecScopeWalkResult = helper.getModule().walk([&](Operation *op) { - if (!VPTOLegalityHelper::isDedicatedVecScopeCarrier(op)) + if (!VPTOLegalityHelper::isDedicatedVecScopeCarrier(op)) { return WalkResult::advance(); + } - if (!VPTOLegalityHelper::getEnclosingVectorScopeCarrier(op)) + if (!VPTOLegalityHelper::getEnclosingVectorScopeCarrier(op)) { return WalkResult::advance(); + } op->emitOpError() << "does not allow nested dedicated pto.vecscope/pto.strict_vecscope"; return WalkResult::interrupt(); }); - if (vecScopeWalkResult.wasInterrupted()) + if (vecScopeWalkResult.wasInterrupted()) { return failure(); + } WalkResult opWalkResult = helper.getModule().walk([&](Operation *op) { (void)VPTOLegalityHelper::inferMaskGranularityFromFamily(op); (void)VPTOLegalityHelper::classifyBufferAddressFamily(op); - if (!VPTOLegalityHelper::requiresVecScope(op)) + if (!VPTOLegalityHelper::requiresVecScope(op)) { return WalkResult::advance(); + } if (VPTOLegalityHelper::getEnclosingVectorScopeCarrier(op)) { if (failed(validateFamilySuffixMaskContracts(op)) || failed(validateUnaryElementTypeContracts(op)) || - failed(validateMaskGranularityContracts(op))) + failed(validateMaskGranularityContracts(op))) { return WalkResult::interrupt(); + } emitHardwareSupportWarnings(op); return WalkResult::advance(); } @@ -1085,16 +1166,18 @@ class VPTOLegalityValidator { FunctionType functionType = func.getFunctionType(); for (auto [idx, inputType] : llvm::enumerate(functionType.getInputs())) { - if (!isa(inputType)) + if (!isa(inputType)) { continue; + } return func.emitError() << "emission-stage VPTO legality rejects memref argument #" << idx << ": " << inputType; } for (auto [idx, resultType] : llvm::enumerate(functionType.getResults())) { - if (!isa(resultType)) + if (!isa(resultType)) { continue; + } return func.emitError() << "emission-stage VPTO legality rejects memref result #" << idx << ": " << resultType; @@ -1119,8 +1202,9 @@ class VPTOLegalityValidator { if (family == VPTOBufferAddressFamily::BufferLike) { for (OpOperand *operand : VPTOLegalityHelper::collectBufferOperands(op)) { Type operandType = operand->get().getType(); - if (!isa(operandType)) + if (!isa(operandType)) { continue; + } op->emitOpError() << "emission-stage VPTO legality rejects memref-form buffer " @@ -1143,8 +1227,9 @@ class VPTOLegalityValidator { } void writeDiagnostic(StringRef message) const { - if (diagOS) + if (diagOS) { *diagOS << message; + } } VPTOLegalityHelper helper; @@ -1168,8 +1253,9 @@ struct PTOValidateVPTOIRPass void runOnOperation() override { ModuleOp module = getOperation(); - if (failed(validateVPTOAuthoringIR(module, &llvm::errs()))) + if (failed(validateVPTOAuthoringIR(module, &llvm::errs()))) { signalPassFailure(); + } } }; @@ -1188,8 +1274,9 @@ struct PTOValidateVPTOEmissionIRPass void runOnOperation() override { ModuleOp module = getOperation(); - if (failed(validateVPTOEmissionIR(module, &llvm::errs()))) + if (failed(validateVPTOEmissionIR(module, &llvm::errs()))) { signalPassFailure(); + } } }; diff --git a/lib/PTO/Transforms/PTOVerifyTFreePass.cpp b/lib/PTO/Transforms/PTOVerifyTFreePass.cpp index fd9c7e1d4b..0bee6073a9 100644 --- a/lib/PTO/Transforms/PTOVerifyTFreePass.cpp +++ b/lib/PTO/Transforms/PTOVerifyTFreePass.cpp @@ -35,8 +35,9 @@ static TFreeOp findMatchingTFree(TPopOp tpopOp) { for (auto it = std::next(tpopOp->getIterator()), end = block->end(); it != end; ++it) { if (auto tfreeOp = dyn_cast(&*it)) { - if (tfreeOp.getPipeHandle() == pipeHandle) + if (tfreeOp.getPipeHandle() == pipeHandle) { return tfreeOp; + } } } return {}; @@ -46,8 +47,9 @@ static Operation *getTopLevelAncestorInBlock(Operation *op, Block *block) { Operation *current = op; while (current && current->getBlock() != block) { Region *parentRegion = current->getParentRegion(); - if (!parentRegion) + if (!parentRegion) { return nullptr; + } current = parentRegion->getParentOp(); } return current; @@ -57,8 +59,9 @@ static bool hasSamePipeTPopInRegion(Operation *op, Value pipeHandle, TPopOp current) { bool found = false; op->walk([&](TPopOp nestedTpop) { - if (nestedTpop == current) + if (nestedTpop == current) { return WalkResult::advance(); + } if (nestedTpop.getPipeHandle() == pipeHandle) { found = true; return WalkResult::interrupt(); @@ -70,8 +73,9 @@ static bool hasSamePipeTPopInRegion(Operation *op, Value pipeHandle, static LogicalResult verifySingleOutstandingUntil(TPopOp tpopOp, Operation *freeBoundary) { - if (!freeBoundary || freeBoundary == tpopOp.getOperation()) + if (!freeBoundary || freeBoundary == tpopOp.getOperation()) { return success(); + } Value pipeHandle = tpopOp.getPipeHandle(); Block *block = tpopOp->getBlock(); @@ -82,8 +86,9 @@ static LogicalResult verifySingleOutstandingUntil(TPopOp tpopOp, return tpopOp.emitOpError( "multiple outstanding pops on the same pipe are not supported"); } - if (op == freeBoundary) + if (op == freeBoundary) { break; + } } return success(); @@ -126,8 +131,9 @@ struct PTOVerifyTFreePass funcOp.walk([&](TPopOp op) { tpops.push_back(op); }); for (TPopOp tpopOp : tpops) { - if (!isInsideSectionOrAttributedKernel(tpopOp, funcOp)) + if (!isInsideSectionOrAttributedKernel(tpopOp, funcOp)) { continue; + } TFreeOp existingTFree = findMatchingTFree(tpopOp); if (!existingTFree) { diff --git a/lib/PTO/Transforms/SIMTPersistentFragmentAnalysis.cpp b/lib/PTO/Transforms/SIMTPersistentFragmentAnalysis.cpp index 32835fc2eb..989a1e2d39 100644 --- a/lib/PTO/Transforms/SIMTPersistentFragmentAnalysis.cpp +++ b/lib/PTO/Transforms/SIMTPersistentFragmentAnalysis.cpp @@ -394,14 +394,16 @@ collectResidentElements(pto::SectionSimtOp initSection, llvm::DenseSet &residentElementSet) { for (Operation &op : initSection.getBody().front()) { auto accessIt = discovery.accessIndices.find(&op); - if (accessIt == discovery.accessIndices.end()) + if (accessIt == discovery.accessIndices.end()) { continue; + } for (unsigned accessIndex : accessIt->second) { const NormalizedPersistentAccess &access = discovery.accesses[accessIndex]; - if (!residentElementSet.insert(access.elementOffset).second) + if (!residentElementSet.insert(access.elementOffset).second) { continue; + } if (!isa(access.op)) { return access.op->emitOpError() @@ -559,8 +561,9 @@ materializeResidentAccessLanes(const PersistentMaterializationPlan &plan, } for (Operation &op : section.getBody().front()) { auto accessIt = discovery.accessIndices.find(&op); - if (accessIt == discovery.accessIndices.end()) + if (accessIt == discovery.accessIndices.end()) { continue; + } for (unsigned accessIndex : accessIt->second) { if (accessIndex >= discovery.accesses.size()) { @@ -713,8 +716,9 @@ SIMTPersistentFragmentAnalysis::SIMTPersistentFragmentAnalysis( func::FuncOp func) { SmallVector persistentAllocas; func.walk([&](LLVM::AllocaOp allocaOp) { - if (allocaOp->hasAttr(kPersistentAttrName)) + if (allocaOp->hasAttr(kPersistentAttrName)) { persistentAllocas.push_back(allocaOp); + } }); // A function without persistent allocations has a valid empty plan and does diff --git a/lib/PTO/Transforms/SlotAffineAnalysis.cpp b/lib/PTO/Transforms/SlotAffineAnalysis.cpp index 23d0e81e38..23cc1f75b6 100644 --- a/lib/PTO/Transforms/SlotAffineAnalysis.cpp +++ b/lib/PTO/Transforms/SlotAffineAnalysis.cpp @@ -60,8 +60,9 @@ struct SlotForm { static bool tryGetConstantInt(Value v, int64_t &out) { IntegerAttr attr; - if (!matchPattern(v, m_Constant(&attr))) + if (!matchPattern(v, m_Constant(&attr))) { return false; + } out = attr.getValue().getSExtValue(); return true; } @@ -71,8 +72,9 @@ static bool tryGetConstantInt(Value v, int64_t &out) { // neither side is a constant. static bool peelAddSubConst(Value v, Value &remaining, int64_t &offset) { Operation *op = v.getDefiningOp(); - if (!op) + if (!op) { return false; + } Value lhs, rhs; bool isSub = false; if (auto add = dyn_cast(op)) { @@ -106,8 +108,9 @@ static bool peelAddSubConst(Value v, Value &remaining, int64_t &offset) { // without a `remui`, treat N as the caller-supplied `expectN` and reduce. // Returns false if the form is not representable. static bool extractSlotForm(Value slot, uint32_t expectN, SlotForm &out) { - if (!slot) + if (!slot) { return false; + } out.innerSym = Value(); out.innerOffset = 0; @@ -118,8 +121,9 @@ static bool extractSlotForm(Value slot, uint32_t expectN, SlotForm &out) { // Case 1: `arith.remui inner, %const_N`. if (auto remOp = dyn_cast_if_present(def)) { int64_t n; - if (!tryGetConstantInt(remOp.getRhs(), n) || n <= 0) + if (!tryGetConstantInt(remOp.getRhs(), n) || n <= 0) { return false; + } out.N = static_cast(n); Value inner = remOp.getLhs(); int64_t offset = 0; @@ -161,8 +165,9 @@ static bool extractSlotForm(Value slot, uint32_t expectN, SlotForm &out) { static int64_t pyMod(int64_t a, int64_t n) { int64_t r = a % n; - if (r < 0) + if (r < 0) { r += n; + } return r; } @@ -203,12 +208,14 @@ SlotRelation compareSlotSSA(Value a, Value b, uint32_t N) { // One side const, other symbolic: cannot prove disjoint without // assuming a value range on the symbol. Equality also unprovable. - if (!fa.innerSym || !fb.innerSym) + if (!fa.innerSym || !fb.innerSym) { return SlotRelation::kUnknown; + } // Both symbolic. Need same symbol to reason about (a - b) mod N. - if (fa.innerSym != fb.innerSym) + if (fa.innerSym != fb.innerSym) { return SlotRelation::kUnknown; + } int64_t diff = pyMod(fa.innerOffset - fb.innerOffset, N); return diff == 0 ? SlotRelation::kEqual : SlotRelation::kDisjoint; diff --git a/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp b/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp index 44f50acb2e..4016137612 100644 --- a/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp +++ b/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp @@ -23,12 +23,15 @@ namespace pto { namespace { static int64_t getConstantIndexOrDynamic(Value value) { - if (!value) + if (!value) { return ShapedType::kDynamic; - if (auto cst = value.getDefiningOp()) + } + if (auto cst = value.getDefiningOp()) { return cst.value(); - if (auto cst = value.getDefiningOp()) + } + if (auto cst = value.getDefiningOp()) { return cst.value(); + } return ShapedType::kDynamic; } @@ -50,14 +53,18 @@ static SmallVector getValidShapeVec(Type type) { /// result, so that two such ops with the same name, attributes and equivalent /// operands are guaranteed to produce the same value. static bool isShapeComputableOp(Operation *op) { - if (!op) + if (!op) { return false; - if (op->getNumRegions() != 0) + } + if (op->getNumRegions() != 0) { return false; - if (op->getNumResults() != 1) + } + if (op->getNumResults() != 1) { return false; - if (!isMemoryEffectFree(op)) + } + if (!isMemoryEffectFree(op)) { return false; + } // Only allow arith ops that appear in typical valid-shape computations // (minsi, maxsi, cmpi, select, addi, subi, muli, divsi, divui, @@ -87,10 +94,12 @@ class StructuralSignatureMap { Attribute attrs; bool operator==(const Key &rhs) const { - if (opNamePtr != rhs.opNamePtr) + if (opNamePtr != rhs.opNamePtr) { return false; - if (operands.size() != rhs.operands.size()) + } + if (operands.size() != rhs.operands.size()) { return false; + } for (auto [l, r] : llvm::zip(operands, rhs.operands)) if (l != r) return false; @@ -133,12 +142,14 @@ class StructuralSignatureMap { static Value canonicalizeValue(Value value, DenseMap &canonicalByValue, StructuralSignatureMap &signatureMap) { - if (!value) + if (!value) { return value; + } auto cachedIt = canonicalByValue.find(value); - if (cachedIt != canonicalByValue.end()) + if (cachedIt != canonicalByValue.end()) { return cachedIt->second; + } // BlockArguments are their own canonical form. if (auto arg = dyn_cast(value)) { @@ -182,7 +193,7 @@ static Value canonicalizeValue(Value value, return representative; } -static constexpr unsigned kInvalidShapeDim = ~0u; +static constexpr unsigned kInvalidShapeDim = ~0U; struct ShapeValueDims { unsigned rows = kInvalidShapeDim; @@ -206,26 +217,31 @@ class ShapeConstraintSolver { unsigned find(unsigned dim) { assert(dim < parent.size() && "shape dim out of range"); - if (parent[dim] == dim) + if (parent[dim] == dim) { return dim; + } parent[dim] = find(parent[dim]); return parent[dim]; } void merge(unsigned lhs, unsigned rhs) { - if (lhs == kInvalidShapeDim || rhs == kInvalidShapeDim) + if (lhs == kInvalidShapeDim || rhs == kInvalidShapeDim) { return; + } unsigned lhsRoot = find(lhs); unsigned rhsRoot = find(rhs); - if (lhsRoot == rhsRoot) + if (lhsRoot == rhsRoot) { return; + } - if (rank[lhsRoot] < rank[rhsRoot]) + if (rank[lhsRoot] < rank[rhsRoot]) { std::swap(lhsRoot, rhsRoot); + } parent[rhsRoot] = lhsRoot; - if (rank[lhsRoot] == rank[rhsRoot]) + if (rank[lhsRoot] == rank[rhsRoot]) { ++rank[lhsRoot]; + } conflicts[lhsRoot] = conflicts[lhsRoot] || conflicts[rhsRoot]; if (constants[lhsRoot] && constants[rhsRoot] && @@ -236,25 +252,29 @@ class ShapeConstraintSolver { } void bindConstant(unsigned dim, int64_t value) { - if (dim == kInvalidShapeDim || value == ShapedType::kDynamic) + if (dim == kInvalidShapeDim || value == ShapedType::kDynamic) { return; + } unsigned root = find(dim); - if (constants[root] && *constants[root] != value) + if (constants[root] && *constants[root] != value) { conflicts[root] = true; - else + } else { constants[root] = value; + } } bool hasConflict(unsigned dim) { - if (dim == kInvalidShapeDim) + if (dim == kInvalidShapeDim) { return true; + } return conflicts[find(dim)]; } std::optional getConstant(unsigned dim) { - if (dim == kInvalidShapeDim) + if (dim == kInvalidShapeDim) { return std::nullopt; + } return constants[find(dim)]; } @@ -262,8 +282,9 @@ class ShapeConstraintSolver { /// buildIterationDomainInfo from proving a consistent shape. Used when a /// runtime set_validshape invalidates alloc-time shape assumptions. void markConflict(unsigned dim) { - if (dim == kInvalidShapeDim) + if (dim == kInvalidShapeDim) { return; + } conflicts[find(dim)] = true; } @@ -279,8 +300,9 @@ static void bindDimToValue(ShapeConstraintSolver &solver, DenseMap &canonicalByValue, StructuralSignatureMap &signatureMap, unsigned dim, Value value) { - if (!value || dim == kInvalidShapeDim) + if (!value || dim == kInvalidShapeDim) { return; + } int64_t constant = getConstantIndexOrDynamic(value); if (constant != ShapedType::kDynamic) { @@ -294,8 +316,9 @@ static void bindDimToValue(ShapeConstraintSolver &solver, canonicalizeValue(value, canonicalByValue, signatureMap); auto [it, inserted] = symbolDimByValue.try_emplace(canonical, kInvalidShapeDim); - if (inserted) + if (inserted) { it->second = solver.createDim(); + } solver.merge(dim, it->second); } @@ -341,18 +364,21 @@ static ShapeValueDims getValueDims( DenseMap &canonicalByValue, StructuralSignatureMap &signatureMap, Value value) { auto existing = dimsByValue.find(value); - if (existing != dimsByValue.end()) + if (existing != dimsByValue.end()) { return existing->second; + } ShapeValueDims dims; SmallVector validShape = getValidShapeVec(value.getType()); if (validShape.size() >= 2) { dims.rows = solver.createDim(); dims.cols = solver.createDim(); - if (!ShapedType::isDynamic(validShape[0])) + if (!ShapedType::isDynamic(validShape[0])) { solver.bindConstant(dims.rows, validShape[0]); - if (!ShapedType::isDynamic(validShape[1])) + } + if (!ShapedType::isDynamic(validShape[1])) { solver.bindConstant(dims.cols, validShape[1]); + } bindExplicitValidDims(solver, symbolDimByValue, canonicalByValue, signatureMap, value, dims); } @@ -382,8 +408,9 @@ static void mergeAllShapes( DenseMap &symbolDimByValue, DenseMap &canonicalByValue, StructuralSignatureMap &signatureMap, ArrayRef values) { - if (values.empty()) + if (values.empty()) { return; + } ShapeValueDims anchor = getValueDims(solver, dimsByValue, symbolDimByValue, canonicalByValue, signatureMap, values.front()); @@ -414,8 +441,9 @@ static void applyShapeConstraintsForNode( signatureMap, semantics.tileOutputs); return; case FusionComputeFamily::RowBroadcastBinary: { - if (semantics.tileOutputs.empty()) + if (semantics.tileOutputs.empty()) { return; + } ShapeValueDims output = getValueDims( solver, dimsByValue, symbolDimByValue, canonicalByValue, signatureMap, semantics.tileOutputs.front()); @@ -442,8 +470,9 @@ static void applyShapeConstraintsForNode( case FusionComputeFamily::ReduceCol: { mergeAllShapes(solver, dimsByValue, symbolDimByValue, canonicalByValue, signatureMap, semantics.tileInputs); - if (semantics.tileInputs.empty() || semantics.tileOutputs.empty()) + if (semantics.tileInputs.empty() || semantics.tileOutputs.empty()) { return; + } ShapeValueDims input = getValueDims( solver, dimsByValue, symbolDimByValue, canonicalByValue, signatureMap, semantics.tileInputs.front()); @@ -504,8 +533,9 @@ static ShapeValueDims getIterationDomainDimsForNode( static IterationDomainInfo buildIterationDomainInfo(ShapeConstraintSolver &solver, ShapeValueDims dims) { IterationDomainInfo info; - if (!dims.isValid()) + if (!dims.isValid()) { return info; + } if (solver.hasConflict(dims.rows) || solver.hasConflict(dims.cols)) { info.unprovenReason = IterationDomainUnprovenReason::InconsistentShape; return info; @@ -513,10 +543,12 @@ buildIterationDomainInfo(ShapeConstraintSolver &solver, ShapeValueDims dims) { info.proof = IterationDomainProof::Proven; info.unprovenReason = IterationDomainUnprovenReason::None; - if (std::optional row = solver.getConstant(dims.rows)) + if (std::optional row = solver.getConstant(dims.rows)) { info.vRow = *row; - if (std::optional col = solver.getConstant(dims.cols)) + } + if (std::optional col = solver.getConstant(dims.cols)) { info.vCol = *col; + } return info; } @@ -572,8 +604,9 @@ struct Rank2IterationSpace { static std::optional getRank2IterationSpace(Value value) { SmallVector validShape = getValidShapeVec(value.getType()); - if (validShape.size() < 2) + if (validShape.size() < 2) { return std::nullopt; + } return Rank2IterationSpace{validShape[0], validShape[1]}; } @@ -581,8 +614,9 @@ static void mergeIterationDim(int64_t &mergedDim, int64_t dim, IterationDomainInfo &info) { if (mergedDim == ShapedType::kDynamic || dim == ShapedType::kDynamic) { mergedDim = ShapedType::kDynamic; - if (info.unprovenReason == IterationDomainUnprovenReason::None) + if (info.unprovenReason == IterationDomainUnprovenReason::None) { info.unprovenReason = IterationDomainUnprovenReason::DynamicShape; + } return; } @@ -597,19 +631,22 @@ inferConsensusIterationDomain(ArrayRef anchorValues) { IterationDomainInfo info; info.unprovenReason = IterationDomainUnprovenReason::None; - if (anchorValues.empty()) + if (anchorValues.empty()) { return info; + } std::optional firstSpace = getRank2IterationSpace(anchorValues.front()); - if (!firstSpace) + if (!firstSpace) { return info; + } info.vRow = firstSpace->rows; info.vCol = firstSpace->cols; - if (info.vRow == ShapedType::kDynamic || info.vCol == ShapedType::kDynamic) + if (info.vRow == ShapedType::kDynamic || info.vCol == ShapedType::kDynamic) { info.unprovenReason = IterationDomainUnprovenReason::DynamicShape; + } for (Value value : ArrayRef(anchorValues).drop_front()) { std::optional space = getRank2IterationSpace(value); @@ -629,8 +666,9 @@ inferConsensusIterationDomain(ArrayRef anchorValues) { return info; } - if (info.unprovenReason == IterationDomainUnprovenReason::None) + if (info.unprovenReason == IterationDomainUnprovenReason::None) { info.unprovenReason = IterationDomainUnprovenReason::DynamicShape; + } return info; } @@ -730,8 +768,9 @@ static LogicalResult inferDynamicIterationDomain(FusionBlockAnalysis &analysis) if (analysis.block) { for (Operation &op : *analysis.block) { auto setVS = dyn_cast(op); - if (!setVS) + if (!setVS) { continue; + } Value source = setVS.getSource(); auto dimsIt = dimsByValue.find(source); if (dimsIt != dimsByValue.end()) { @@ -780,10 +819,12 @@ static Value getWriteInstanceStorageValue(Operation *op, unsigned outputIndex, if (auto dpsIface = dyn_cast(op)) { unsigned tileOutputIndex = 0; for (Value init : dpsIface.getDpsInits()) { - if (!isa(init.getType())) + if (!isa(init.getType())) { continue; - if (tileOutputIndex == outputIndex) + } + if (tileOutputIndex == outputIndex) { return init; + } ++tileOutputIndex; } } @@ -803,14 +844,16 @@ static unsigned getOrCreateLivenessSlot(DenseMap &slotByValue, } static void appendUniqueNode(SmallVectorImpl &nodes, unsigned nodeId) { - if (!llvm::is_contained(nodes, nodeId)) + if (!llvm::is_contained(nodes, nodeId)) { nodes.push_back(nodeId); + } } static void recordLastLocalConsumer(std::optional &lastLocalConsumer, unsigned consumerId) { - if (!lastLocalConsumer || consumerId > *lastLocalConsumer) + if (!lastLocalConsumer || consumerId > *lastLocalConsumer) { lastLocalConsumer = consumerId; + } } static void finalizeBlockLiveness( @@ -827,17 +870,20 @@ static void finalizeBlockLiveness( } auto kindIt = kindByOp.find(user); - if (kindIt == kindByOp.end()) + if (kindIt == kindByOp.end()) { continue; + } - if (user->hasTrait()) + if (user->hasTrait()) { state.live.escapesBlock = true; + } switch (kindIt->second) { case FusionOpKind::Compute: { auto nodeIt = computeNodeByOp.find(user); - if (nodeIt == computeNodeByOp.end()) + if (nodeIt == computeNodeByOp.end()) { continue; + } unsigned consumerId = nodeIt->second; appendUniqueNode(state.live.consumerNodes, consumerId); recordLastLocalConsumer(state.live.lastLocalConsumer, consumerId); @@ -858,11 +904,13 @@ static std::optional findReachingWriteInstance( ArrayRef writeInstanceIds, ArrayRef mutableWriteInstances, std::optional userBlockOrder) { - if (writeInstanceIds.empty()) + if (writeInstanceIds.empty()) { return std::nullopt; + } - if (!userBlockOrder) + if (!userBlockOrder) { return writeInstanceIds.back(); + } for (unsigned writeInstanceId : llvm::reverse(writeInstanceIds)) { if (mutableWriteInstances[writeInstanceId].producerBlockOrder < @@ -874,8 +922,9 @@ static std::optional findReachingWriteInstance( static bool isDpsInitOperandUse(OpOperand &use) { auto dpsIface = dyn_cast(use.getOwner()); - if (!dpsIface) + if (!dpsIface) { return false; + } for (OpOperand &dpsInit : dpsIface.getDpsInitsMutable()) if (&dpsInit == &use) @@ -890,27 +939,31 @@ static void finalizeWriteInstances( ArrayRef mutableLiveness, SmallVectorImpl &mutableWriteInstances) { for (const MutableLiveness &storageState : mutableLiveness) { - if (storageState.live.writeInstances.empty()) + if (storageState.live.writeInstances.empty()) { continue; + } for (OpOperand &use : storageState.live.value.getUses()) { - if (isDpsInitOperandUse(use)) + if (isDpsInitOperandUse(use)) { continue; + } Operation *user = use.getOwner(); bool isInBlock = user->getBlock() == █ std::optional userBlockOrder; if (isInBlock) { auto orderIt = blockOrderByOp.find(user); - if (orderIt != blockOrderByOp.end()) + if (orderIt != blockOrderByOp.end()) { userBlockOrder = orderIt->second; + } } std::optional writeInstanceId = findReachingWriteInstance( storageState.live.writeInstances, mutableWriteInstances, userBlockOrder); - if (!writeInstanceId) + if (!writeInstanceId) { continue; + } FusionWriteInstanceLiveness &writeLive = mutableWriteInstances[*writeInstanceId].live; @@ -922,17 +975,20 @@ static void finalizeWriteInstances( } auto kindIt = kindByOp.find(user); - if (kindIt == kindByOp.end()) + if (kindIt == kindByOp.end()) { continue; + } - if (user->hasTrait()) + if (user->hasTrait()) { writeLive.escapesBlock = true; + } switch (kindIt->second) { case FusionOpKind::Compute: { auto nodeIt = computeNodeByOp.find(user); - if (nodeIt == computeNodeByOp.end()) + if (nodeIt == computeNodeByOp.end()) { continue; + } unsigned consumerId = nodeIt->second; appendUniqueNode(writeLive.consumerNodes, consumerId); recordLastLocalConsumer(writeLive.lastLocalConsumer, consumerId); @@ -980,10 +1036,12 @@ static FailureOr analyzeBlockDFG(Block &block) { kindByOp[&op] = semanticsOr->kind; if (semanticsOr->kind == FusionOpKind::LocalBoundary) { - for (Value input : semanticsOr->tileInputs) + for (Value input : semanticsOr->tileInputs) { getOrCreateLivenessSlot(livenessSlotByValue, mutableLiveness, input); - for (Value output : semanticsOr->tileOutputs) + } + for (Value output : semanticsOr->tileOutputs) { getOrCreateLivenessSlot(livenessSlotByValue, mutableLiveness, output); + } ++blockOrder; continue; } @@ -1026,8 +1084,9 @@ static FailureOr analyzeBlockDFG(Block &block) { node.id); auto producerIt = producerByValue.find(input); - if (producerIt == producerByValue.end()) + if (producerIt == producerByValue.end()) { continue; + } FusionDFGEdge edge; edge.producerNode = producerIt->second; @@ -1037,8 +1096,9 @@ static FailureOr analyzeBlockDFG(Block &block) { unsigned edgeId = analysis.edges.size(); analysis.edges.push_back(edge); node.incomingEdges.push_back(edgeId); - if (edge.producerNode < analysis.computeNodes.size()) + if (edge.producerNode < analysis.computeNodes.size()) { analysis.computeNodes[edge.producerNode].outgoingEdges.push_back(edgeId); + } } analysis.computeNodes.push_back(std::move(node)); @@ -1050,11 +1110,13 @@ static FailureOr analyzeBlockDFG(Block &block) { mutableLiveness, mutableWriteInstances); analysis.liveness.reserve(mutableLiveness.size()); - for (MutableLiveness &state : mutableLiveness) + for (MutableLiveness &state : mutableLiveness) { analysis.liveness.push_back(std::move(state.live)); + } analysis.writeInstances.reserve(mutableWriteInstances.size()); - for (MutableWriteInstance &state : mutableWriteInstances) + for (MutableWriteInstance &state : mutableWriteInstances) { analysis.writeInstances.push_back(std::move(state.live)); + } return std::move(analysis); } @@ -1063,8 +1125,9 @@ static LogicalResult analyzeRegionDFG(Region ®ion, SmallVectorImpl &blocks) { for (Block &block : region.getBlocks()) { FailureOr blockAnalysis = analyzeBlockDFG(block); - if (failed(blockAnalysis)) + if (failed(blockAnalysis)) { return failure(); + } blocks.push_back(std::move(*blockAnalysis)); for (Operation &op : block) for (Region &nested : op.getRegions()) @@ -1079,8 +1142,9 @@ static LogicalResult analyzeRegionDFG(Region ®ion, FailureOr buildPreFusionAnalysisDFG(func::FuncOp func) { PreFusionAnalysisResult result; - if (failed(analyzeRegionDFG(func.getRegion(), result.blocks))) + if (failed(analyzeRegionDFG(func.getRegion(), result.blocks))) { return failure(); + } return std::move(result); } @@ -1088,11 +1152,13 @@ LogicalResult inferIterationDomainClasses(PreFusionAnalysisResult &result, bool enableShapeInference) { for (FusionBlockAnalysis &block : result.blocks) { if (enableShapeInference) { - if (failed(inferDynamicIterationDomain(block))) + if (failed(inferDynamicIterationDomain(block))) { return failure(); + } } else { - if (failed(inferStaticIterationDomain(block))) + if (failed(inferStaticIterationDomain(block))) { return failure(); + } } } return success(); @@ -1101,10 +1167,12 @@ LogicalResult inferIterationDomainClasses(PreFusionAnalysisResult &result, FailureOr buildPreFusionAnalysis(func::FuncOp func, bool enableShapeInference) { FailureOr result = buildPreFusionAnalysisDFG(func); - if (failed(result)) + if (failed(result)) { return failure(); - if (failed(inferIterationDomainClasses(*result, enableShapeInference))) + } + if (failed(inferIterationDomainClasses(*result, enableShapeInference))) { return failure(); + } return std::move(*result); } diff --git a/lib/PTO/Transforms/TileFusion/FusionOpSemantics.cpp b/lib/PTO/Transforms/TileFusion/FusionOpSemantics.cpp index a9a13eae14..d105485234 100644 --- a/lib/PTO/Transforms/TileFusion/FusionOpSemantics.cpp +++ b/lib/PTO/Transforms/TileFusion/FusionOpSemantics.cpp @@ -42,8 +42,9 @@ static SmallVector collectNormalizedTileOutputs(Operation *op) { if (auto dpsIface = dyn_cast(op)) { for (Value init : dpsIface.getDpsInits()) { - if (isTileFusionTileValue(init)) + if (isTileFusionTileValue(init)) { outputs.push_back(init); + } } if (!outputs.empty()) return outputs; @@ -103,10 +104,11 @@ FailureOr getFusionOpSemantics(Operation *op) { continue; Value value = operand.get(); - if (isTileFusionTileValue(value)) + if (isTileFusionTileValue(value)) { semantics.tileInputs.push_back(value); - else + } else { semantics.scalarInputs.push_back(value); + } } if (semantics.tileInputs.empty()) { diff --git a/lib/PTO/Transforms/TileFusion/PTOFlattenFusionRegion.cpp b/lib/PTO/Transforms/TileFusion/PTOFlattenFusionRegion.cpp index 6f8115d4c0..81db720f1f 100644 --- a/lib/PTO/Transforms/TileFusion/PTOFlattenFusionRegion.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOFlattenFusionRegion.cpp @@ -28,8 +28,9 @@ namespace { static LogicalResult flattenFusionRegion(pto::FusionRegionOp fusionRegion) { Block &body = fusionRegion.getBody().front(); auto yieldOp = dyn_cast(body.getTerminator()); - if (!yieldOp) + if (!yieldOp) { return fusionRegion.emitOpError("expects body to terminate with pto.yield"); + } SmallVector yieldedValues(yieldOp.getValues().begin(), yieldOp.getValues().end()); @@ -60,8 +61,9 @@ struct PTOFlattenFusionRegionPass void runOnOperation() override { func::FuncOp func = getOperation(); - if (func.isExternal()) + if (func.isExternal()) { return; + } SmallVector fusionRegions; func.walk([&](pto::FusionRegionOp fusionRegion) { diff --git a/lib/PTO/Transforms/TileFusion/PTOFusionLoadStoreElision.cpp b/lib/PTO/Transforms/TileFusion/PTOFusionLoadStoreElision.cpp index 9c52282a85..f0adf27c43 100644 --- a/lib/PTO/Transforms/TileFusion/PTOFusionLoadStoreElision.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOFusionLoadStoreElision.cpp @@ -55,20 +55,27 @@ static bool areEquivalentValueRanges(ArrayRef lhs, ArrayRef rhs) { } static bool areEquivalentOperations(Operation *lhs, Operation *rhs) { - if (!lhs || !rhs) + if (!lhs || !rhs) { return false; - if (lhs->getName() != rhs->getName()) + } + if (lhs->getName() != rhs->getName()) { return false; - if (lhs->getNumRegions() != 0 || rhs->getNumRegions() != 0) + } + if (lhs->getNumRegions() != 0 || rhs->getNumRegions() != 0) { return false; - if (lhs->getNumResults() != rhs->getNumResults()) + } + if (lhs->getNumResults() != rhs->getNumResults()) { return false; - if (lhs->getNumOperands() != rhs->getNumOperands()) + } + if (lhs->getNumOperands() != rhs->getNumOperands()) { return false; - if (lhs->getAttrDictionary() != rhs->getAttrDictionary()) + } + if (lhs->getAttrDictionary() != rhs->getAttrDictionary()) { return false; - if (!llvm::equal(lhs->getResultTypes(), rhs->getResultTypes())) + } + if (!llvm::equal(lhs->getResultTypes(), rhs->getResultTypes())) { return false; + } if (auto lhsDim = dyn_cast(lhs)) { auto rhsDim = cast(rhs); @@ -78,19 +85,23 @@ static bool areEquivalentOperations(Operation *lhs, Operation *rhs) { for (auto [lhsOperand, rhsOperand] : llvm::zip(lhs->getOperands(), rhs->getOperands())) { - if (!areEquivalentValues(lhsOperand, rhsOperand)) + if (!areEquivalentValues(lhsOperand, rhsOperand)) { return false; + } } return true; } static bool areEquivalentValues(Value lhs, Value rhs) { - if (lhs == rhs) + if (lhs == rhs) { return true; - if (!lhs || !rhs) + } + if (!lhs || !rhs) { return false; - if (lhs.getType() != rhs.getType()) + } + if (lhs.getType() != rhs.getType()) { return false; + } auto lhsArg = dyn_cast(lhs); auto rhsArg = dyn_cast(rhs); @@ -111,22 +122,25 @@ static bool isPureNoRegionOp(Operation *op) { } static bool isSupportedLoopPreludeOp(Operation *op) { - if (isa(op)) + if (isa(op)) { return true; + } return isPureNoRegionOp(op); } static bool isSupportedLeafOp(Operation *op) { - if (isa(op)) + if (isa(op)) { return true; + } return isPureNoRegionOp(op); } static Value getCanonicalTrackedValue(Value value) { while (value) { Operation *def = value.getDefiningOp(); - if (!def) + if (!def) { break; + } if (auto tileBufAddr = dyn_cast(def)) { value = tileBufAddr.getSrc(); @@ -177,8 +191,9 @@ static Value getCanonicalTrackedValue(Value value) { continue; } if (auto cast = dyn_cast(def)) { - if (cast.getInputs().empty()) + if (cast.getInputs().empty()) { break; + } if (auto result = dyn_cast(value)) { unsigned resultNumber = result.getResultNumber(); if (resultNumber < cast.getInputs().size()) { @@ -206,27 +221,32 @@ static Operation *getTopLevelAncestorInBlock(Operation *op, Block *block) { static Region *getDirectRegionUnderAncestor(Operation *op, Operation *ancestor) { for (Operation *cur = op; cur; cur = cur->getParentOp()) { Operation *parent = cur->getParentOp(); - if (parent == ancestor) + if (parent == ancestor) { return cur->getBlock() ? cur->getBlock()->getParent() : nullptr; + } } return nullptr; } static bool areMutuallyExclusiveByIfRegion(Operation *lhs, Operation *rhs) { - if (!lhs || !rhs) + if (!lhs || !rhs) { return false; + } for (Operation *ancestor = lhs; ancestor; ancestor = ancestor->getParentOp()) { auto ifOp = dyn_cast(ancestor); - if (!ifOp) + if (!ifOp) { continue; + } Region *lhsRegion = getDirectRegionUnderAncestor(lhs, ifOp); Region *rhsRegion = getDirectRegionUnderAncestor(rhs, ifOp); - if (!lhsRegion || !rhsRegion) + if (!lhsRegion || !rhsRegion) { continue; - if (lhsRegion != rhsRegion) + } + if (lhsRegion != rhsRegion) { return true; + } } return false; @@ -236,8 +256,9 @@ static std::optional buildFusionRegionStoreContext(pto::FusionRegionOp fusionRegion) { Block &body = fusionRegion.getBody().front(); auto yieldOp = dyn_cast(body.getTerminator()); - if (!yieldOp) + if (!yieldOp) { return std::nullopt; + } FusionRegionStoreContext context; context.body = &body; @@ -246,23 +267,26 @@ buildFusionRegionStoreContext(pto::FusionRegionOp fusionRegion) { for (Value yielded : yieldOp.getValues()) { Value canonical = getCanonicalTrackedValue(yielded); - if (canonical) + if (canonical) { context.yieldedValues.insert(canonical); + } } return context; } static bool isSupportedLoopRoot(scf::ForOp loop) { - if (!loop) + if (!loop) { return false; + } return isa( loop->getParentOp()); } static Block *getLeafLoopBody(scf::ForOp carrierLoop) { - if (!carrierLoop) + if (!carrierLoop) { return nullptr; + } scf::ForOp currentLoop = carrierLoop; while (currentLoop) { @@ -271,16 +295,18 @@ static Block *getLeafLoopBody(scf::ForOp carrierLoop) { for (Operation &op : currentLoop.getBody()->without_terminator()) { bodyOps.push_back(&op); if (auto loop = dyn_cast(op)) { - if (innerLoop) + if (innerLoop) { return nullptr; + } innerLoop = loop; } } if (!innerLoop) { Block *leafBody = currentLoop.getBody(); - if (!leafBody) + if (!leafBody) { return nullptr; + } for (Operation &op : leafBody->without_terminator()) if (!isSupportedLeafOp(&op)) return nullptr; @@ -293,8 +319,9 @@ static Block *getLeafLoopBody(scf::ForOp carrierLoop) { seenInnerLoop = true; continue; } - if (seenInnerLoop || !isSupportedLoopPreludeOp(op)) + if (seenInnerLoop || !isSupportedLoopPreludeOp(op)) { return nullptr; + } } currentLoop = innerLoop; @@ -314,26 +341,31 @@ static Value inferVPTOLoadUserMask(pto::VldsOp load) { Value inferredMask; for (OpOperand &use : load.getResult().getUses()) { Operation *owner = use.getOwner(); - if (!owner || owner->getNumRegions() != 0) + if (!owner || owner->getNumRegions() != 0) { return Value(); + } Value ownerMask; for (Value operand : owner->getOperands()) { - if (!isa(operand.getType())) + if (!isa(operand.getType())) { continue; - if (!ownerMask) + } + if (!ownerMask) { ownerMask = operand; - else if (!areEquivalentMaskValues(ownerMask, operand)) + } else if (!areEquivalentMaskValues(ownerMask, operand)) { return Value(); + } } - if (!ownerMask) + if (!ownerMask) { return Value(); + } - if (!inferredMask) + if (!inferredMask) { inferredMask = ownerMask; - else if (!areEquivalentMaskValues(inferredMask, ownerMask)) + } else if (!areEquivalentMaskValues(inferredMask, ownerMask)) { return Value(); + } } return inferredMask; } @@ -367,20 +399,24 @@ static bool shouldElideTailStore( Operation *scopeOp, const llvm::SmallPtrSetImpl &scheduledForErase) { Value canonicalBase = getCanonicalTrackedValue(store.base); - if (!canonicalBase) + if (!canonicalBase) { return false; + } Operation *localScopeOp = scopeOp ? scopeOp : store.op; - if (!localScopeOp) + if (!localScopeOp) { return false; + } // Yielded frontier is still region-observable in v1, so its final // materializing store must be preserved even if there is no reload. - if (context.yieldedValues.contains(canonicalBase)) + if (context.yieldedValues.contains(canonicalBase)) { return false; + } for (OpOperand &use : canonicalBase.getUses()) { Operation *owner = use.getOwner(); - if (!owner || scheduledForErase.contains(owner)) + if (!owner || scheduledForErase.contains(owner)) { continue; + } if (context.regionOp->isProperAncestor(owner)) { // Uses nested under the current carrier loop are fine: erasing the tail // store only affects memory materialization, while SSA users still @@ -388,15 +424,19 @@ static bool shouldElideTailStore( // fusion region may still require the buffer to stay materialized, so // keep the store. Operation *topLevelUser = getTopLevelAncestorInBlock(owner, context.body); - if (!topLevelUser) + if (!topLevelUser) { return false; - if (scheduledForErase.contains(topLevelUser)) + } + if (scheduledForErase.contains(topLevelUser)) { continue; - if (topLevelUser == localScopeOp) + } + if (topLevelUser == localScopeOp) { continue; + } if (localScopeOp->getBlock() == topLevelUser->getBlock() && - localScopeOp->isBeforeInBlock(topLevelUser)) + localScopeOp->isBeforeInBlock(topLevelUser)) { return false; + } continue; } @@ -405,16 +445,20 @@ static bool shouldElideTailStore( Operation *topLevelUser = getTopLevelAncestorInBlock(owner, context.parentBlock); if (!topLevelUser) { - if (areMutuallyExclusiveByIfRegion(localScopeOp, owner)) + if (areMutuallyExclusiveByIfRegion(localScopeOp, owner)) { continue; + } return false; } - if (scheduledForErase.contains(topLevelUser)) + if (scheduledForErase.contains(topLevelUser)) { continue; - if (topLevelUser == context.regionOp) + } + if (topLevelUser == context.regionOp) { continue; - if (context.regionOp->isBeforeInBlock(topLevelUser)) + } + if (context.regionOp->isBeforeInBlock(topLevelUser)) { return false; + } } return true; } @@ -480,21 +524,24 @@ static bool elideLoadStoreRoundTripsInLeafBody( continue; } - if (!isPureNoRegionOp(&op)) + if (!isPureNoRegionOp(&op)) { trackedStores.clear(); + } } if (context) { for (const TrackedStore &store : trackedStores) { - if (!shouldElideTailStore(store, *context, scopeOp, scheduledForErase)) + if (!shouldElideTailStore(store, *context, scopeOp, scheduledForErase)) { continue; + } scheduleErase(store.op); changed = true; } } - for (Operation *op : eraseOrder) + for (Operation *op : eraseOrder) { op->erase(); + } return changed; } @@ -506,8 +553,9 @@ struct PTOFusionLoadStoreElisionPass void runOnOperation() override { func::FuncOp func = getOperation(); - if (func.isExternal()) + if (func.isExternal()) { return; + } bool changed = false; @@ -515,32 +563,37 @@ struct PTOFusionLoadStoreElisionPass func.walk([&](pto::FusionRegionOp fusionRegion) { std::optional context = buildFusionRegionStoreContext(fusionRegion); - if (!context) + if (!context) { return; + } regionContexts.try_emplace(fusionRegion.getOperation(), std::move(*context)); }); func.walk([&](pto::FusionRegionOp fusionRegion) { auto it = regionContexts.find(fusionRegion.getOperation()); - if (it == regionContexts.end()) + if (it == regionContexts.end()) { return; + } Block &body = fusionRegion.getBody().front(); - if (!isSupportedStraightLineBlock(body)) + if (!isSupportedStraightLineBlock(body)) { return; + } changed |= elideLoadStoreRoundTripsInLeafBody(body, &it->second, nullptr); }); auto runElisionForLeafBody = [&](Block *leafBody, Operation *scopeOp, pto::FusionRegionOp fusionRegion) { - if (!leafBody || !fusionRegion) + if (!leafBody || !fusionRegion) { return; + } auto it = regionContexts.find(fusionRegion.getOperation()); - if (it == regionContexts.end()) + if (it == regionContexts.end()) { return; + } changed |= elideLoadStoreRoundTripsInLeafBody(*leafBody, &it->second, scopeOp); @@ -548,28 +601,32 @@ struct PTOFusionLoadStoreElisionPass func.walk([&](pto::VecScopeOp vecscope) { if (auto fusionRegion = vecscope->getParentOfType()) { - if (isSupportedStraightLineBlock(vecscope.getBody().front())) + if (isSupportedStraightLineBlock(vecscope.getBody().front())) { runElisionForLeafBody(&vecscope.getBody().front(), vecscope, fusionRegion); + } } }); func.walk([&](pto::StrictVecScopeOp vecscope) { if (auto fusionRegion = vecscope->getParentOfType()) { - if (isSupportedStraightLineBlock(vecscope.getBody().front())) + if (isSupportedStraightLineBlock(vecscope.getBody().front())) { runElisionForLeafBody(&vecscope.getBody().front(), vecscope, fusionRegion); + } } }); func.walk([&](scf::ForOp loop) { - if (!isSupportedLoopRoot(loop)) + if (!isSupportedLoopRoot(loop)) { return; + } runElisionForLeafBody(getLeafLoopBody(loop), loop.getOperation(), loop->getParentOfType()); }); - if (!changed) + if (!changed) { markAllAnalysesPreserved(); + } } }; diff --git a/lib/PTO/Transforms/TileFusion/PTOFusionPlan.cpp b/lib/PTO/Transforms/TileFusion/PTOFusionPlan.cpp index 5d936b8f5d..e29e888b02 100644 --- a/lib/PTO/Transforms/TileFusion/PTOFusionPlan.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOFusionPlan.cpp @@ -91,8 +91,9 @@ static bool isCurrentlyPlannableOp(StringRef opName) { static bool isProvenIterationDomain( const pto::FusionBlockAnalysis &blockAnalysis, const pto::FusionComputeNode &node) { - if (node.iterationDomainClass >= blockAnalysis.iterationDomainClasses.size()) + if (node.iterationDomainClass >= blockAnalysis.iterationDomainClasses.size()) { return false; + } return blockAnalysis.iterationDomainClasses[node.iterationDomainClass] .info.proof == pto::IterationDomainProof::Proven; } @@ -128,10 +129,12 @@ static bool dependsOnPreviousNode( const pto::FusionComputeNode &previous, const pto::FusionComputeNode ¤t) { for (unsigned edgeId : current.incomingEdges) { - if (edgeId >= blockAnalysis.edges.size()) + if (edgeId >= blockAnalysis.edges.size()) { continue; - if (blockAnalysis.edges[edgeId].producerNode == previous.id) + } + if (blockAnalysis.edges[edgeId].producerNode == previous.id) { return true; + } } for (Value output : previous.semantics.tileOutputs) @@ -147,8 +150,9 @@ buildStableInGroupOrder(ArrayRef members) { members.end()); llvm::stable_sort(ordered, [](const pto::FusionComputeNode *lhs, const pto::FusionComputeNode *rhs) { - if (lhs->blockOrder != rhs->blockOrder) + if (lhs->blockOrder != rhs->blockOrder) { return lhs->blockOrder < rhs->blockOrder; + } return lhs->id < rhs->id; }); return ordered; @@ -159,15 +163,17 @@ static void assignStableGroupMetadata(ArrayRef groups, int64_t &nextGroupId) { SmallVector orderedGroups; orderedGroups.reserve(groups.size()); - for (const PlannedFusionGroup &group : groups) + for (const PlannedFusionGroup &group : groups) { orderedGroups.push_back(&group); + } llvm::stable_sort(orderedGroups, [](const PlannedFusionGroup *lhs, const PlannedFusionGroup *rhs) { const pto::FusionComputeNode *lhsFirst = lhs->members.front(); const pto::FusionComputeNode *rhsFirst = rhs->members.front(); - if (lhsFirst->blockOrder != rhsFirst->blockOrder) + if (lhsFirst->blockOrder != rhsFirst->blockOrder) { return lhsFirst->blockOrder < rhsFirst->blockOrder; + } return lhsFirst->id < rhsFirst->id; }); @@ -201,10 +207,12 @@ countEdgesFromGroup(const pto::FusionBlockAnalysis &blockAnalysis, unsigned count = 0; for (unsigned edgeId : candidate.incomingEdges) { - if (edgeId >= blockAnalysis.edges.size()) + if (edgeId >= blockAnalysis.edges.size()) { continue; - if (producerIds.contains(blockAnalysis.edges[edgeId].producerNode)) + } + if (producerIds.contains(blockAnalysis.edges[edgeId].producerNode)) { ++count; + } } return count; } @@ -218,17 +226,21 @@ static bool nodesHaveDirectDataFlowConnection( const pto::FusionBlockAnalysis &blockAnalysis, const pto::FusionComputeNode &lhs, const pto::FusionComputeNode &rhs) { for (unsigned edgeId : lhs.outgoingEdges) { - if (edgeId >= blockAnalysis.edges.size()) + if (edgeId >= blockAnalysis.edges.size()) { continue; - if (blockAnalysis.edges[edgeId].consumerNode == rhs.id) + } + if (blockAnalysis.edges[edgeId].consumerNode == rhs.id) { return true; + } } for (unsigned edgeId : lhs.incomingEdges) { - if (edgeId >= blockAnalysis.edges.size()) + if (edgeId >= blockAnalysis.edges.size()) { continue; - if (blockAnalysis.edges[edgeId].producerNode == rhs.id) + } + if (blockAnalysis.edges[edgeId].producerNode == rhs.id) { return true; + } } for (Value output : lhs.semantics.tileOutputs) @@ -269,8 +281,9 @@ computeGroupFootprint(ArrayRef members) { for (const pto::FusionComputeNode *member : members) { for (Value input : member->semantics.tileInputs) { touchedTiles.insert(input); - if (!producedTiles.contains(input)) + if (!producedTiles.contains(input)) { externalInputs.insert(input); + } } } @@ -300,8 +313,9 @@ class ConservativeGreedyCostModel final : public CostModel { evaluateSeed(const PlanningContext &ctx, const pto::FusionComputeNode &candidate) const override { PlanningDecision decision; - if (!isSupportedPlanningNode(candidate)) + if (!isSupportedPlanningNode(candidate)) { return decision; + } if (!isProvenIterationDomain(ctx.blockAnalysis, candidate)) { decision.cost.rejectedForDynamicShape = true; @@ -317,8 +331,9 @@ class ConservativeGreedyCostModel final : public CostModel { ArrayRef currentGroup, const pto::FusionComputeNode &candidate) const override { PlanningDecision seedDecision = evaluateSeed(ctx, candidate); - if (!seedDecision.accept) + if (!seedDecision.accept) { return seedDecision; + } PlanningDecision decision; if (currentGroup.empty()) { @@ -333,8 +348,9 @@ class ConservativeGreedyCostModel final : public CostModel { candidate.blockOrder == previous.blockOrder + 1; const bool directlyDependent = dependsOnPreviousNode(ctx.blockAnalysis, previous, candidate); - if (!sameDomainClass || !contiguousInBlock || !directlyDependent) + if (!sameDomainClass || !contiguousInBlock || !directlyDependent) { return decision; + } SmallVector proposedGroup( currentGroup.begin(), currentGroup.end()); @@ -360,8 +376,9 @@ class ConservativeDAGGreedyCostModel final : public CostModel { evaluateSeed(const PlanningContext &ctx, const pto::FusionComputeNode &candidate) const override { PlanningDecision decision; - if (!isSupportedPlanningNode(candidate)) + if (!isSupportedPlanningNode(candidate)) { return decision; + } if (!isProvenIterationDomain(ctx.blockAnalysis, candidate)) { decision.cost.rejectedForDynamicShape = true; @@ -377,8 +394,9 @@ class ConservativeDAGGreedyCostModel final : public CostModel { ArrayRef currentGroup, const pto::FusionComputeNode &candidate) const override { PlanningDecision seedDecision = evaluateSeed(ctx, candidate); - if (!seedDecision.accept) + if (!seedDecision.accept) { return seedDecision; + } PlanningDecision decision; if (currentGroup.empty()) { @@ -390,13 +408,15 @@ class ConservativeDAGGreedyCostModel final : public CostModel { candidate.iterationDomainClass) return decision; - if (hasHardBoundaryToGroup(currentGroup, candidate)) + if (hasHardBoundaryToGroup(currentGroup, candidate)) { return decision; + } const unsigned connectionCount = countConnectionsToGroup(ctx.blockAnalysis, currentGroup, candidate); - if (connectionCount == 0) + if (connectionCount == 0) { return decision; + } SmallVector proposedGroup( currentGroup.begin(), currentGroup.end()); @@ -479,12 +499,14 @@ class ConservativeDAGGreedyStrategyEngine final : public StrategyEngine { DenseSet assignedNodes; for (const pto::FusionComputeNode &seed : ctx.blockAnalysis.computeNodes) { - if (assignedNodes.contains(seed.id)) + if (assignedNodes.contains(seed.id)) { continue; + } PlanningDecision seedDecision = costModel.evaluateSeed(ctx, seed); - if (!seedDecision.accept) + if (!seedDecision.accept) { continue; + } SmallVector groupMembers; DenseSet groupNodeIds; @@ -502,8 +524,9 @@ class ConservativeDAGGreedyStrategyEngine final : public StrategyEngine { PlanningDecision appendDecision = costModel.evaluateAppend(ctx, groupMembers, candidate); - if (!appendDecision.accept) + if (!appendDecision.accept) { continue; + } groupMembers.push_back(&candidate); groupNodeIds.insert(candidate.id); @@ -511,14 +534,16 @@ class ConservativeDAGGreedyStrategyEngine final : public StrategyEngine { } } - if (groupMembers.size() < 2) + if (groupMembers.size() < 2) { continue; + } PlannedFusionGroup group; group.members = buildStableInGroupOrder(groupMembers); groups.push_back(group); - for (const pto::FusionComputeNode *member : group.members) + for (const pto::FusionComputeNode *member : group.members) { assignedNodes.insert(member->id); + } } return groups; @@ -554,8 +579,9 @@ struct FusionPlanPass : public pto::impl::FusionPlanBase { void runOnOperation() override { func::FuncOp func = getOperation(); - if (func.isExternal()) + if (func.isExternal()) { return; + } clearPlanningAttrs(func); diff --git a/lib/PTO/Transforms/TileFusion/PTOFusionPredicateElision.cpp b/lib/PTO/Transforms/TileFusion/PTOFusionPredicateElision.cpp index 59ce4d1b60..dc948855ee 100644 --- a/lib/PTO/Transforms/TileFusion/PTOFusionPredicateElision.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOFusionPredicateElision.cpp @@ -122,17 +122,20 @@ static std::optional buildPltCandidate(Operation *op) { static std::optional getForIterArgInfo(Value value) { auto arg = dyn_cast(value); - if (!arg || arg.getArgNumber() == 0) + if (!arg || arg.getArgNumber() == 0) { return std::nullopt; + } auto forOp = dyn_cast_or_null(arg.getParentRegion()->getParentOp()); - if (!forOp || arg.getOwner() != forOp.getBody()) + if (!forOp || arg.getOwner() != forOp.getBody()) { return std::nullopt; + } unsigned iterArgIndex = arg.getArgNumber() - 1; - if (iterArgIndex >= forOp.getInitArgs().size()) + if (iterArgIndex >= forOp.getInitArgs().size()) { return std::nullopt; + } return ForIterArgInfo{forOp, iterArgIndex}; } @@ -140,12 +143,14 @@ static std::optional getPltScalarOutInfo(Value value) { // Element-wise templates keep the remaining element count as index while // plt consumes and returns an integer scalar. Look through the cast used to // feed the scalar result back to scf.for. - if (auto indexCast = value.getDefiningOp()) + if (auto indexCast = value.getDefiningOp()) { value = indexCast.getIn(); + } auto result = dyn_cast(value); - if (!result || result.getResultNumber() != 1) + if (!result || result.getResultNumber() != 1) { return std::nullopt; + } if (auto plt = dyn_cast(result.getOwner())) return PltScalarOutInfo{plt.getScalar(), 8}; @@ -160,10 +165,12 @@ static bool areEquivalentLoopCarriedValues(Value lhs, Value rhs, ValueEquivalenceContext &context) { std::optional lhsInfo = getForIterArgInfo(lhs); std::optional rhsInfo = getForIterArgInfo(rhs); - if (!lhsInfo || !rhsInfo) + if (!lhsInfo || !rhsInfo) { return false; - if (lhsInfo->forOp != rhsInfo->forOp) + } + if (lhsInfo->forOp != rhsInfo->forOp) { return false; + } if (lhsInfo->forOp.getRegionIterArgs().size() != lhsInfo->forOp.getInitArgs().size()) @@ -182,24 +189,29 @@ static bool areEquivalentLoopCarriedValues(Value lhs, Value rhs, getPltScalarOutInfo(yieldedValues[lhsInfo->iterArgIndex]); std::optional rhsYieldInfo = getPltScalarOutInfo(yieldedValues[rhsInfo->iterArgIndex]); - if (!lhsYieldInfo || !rhsYieldInfo) + if (!lhsYieldInfo || !rhsYieldInfo) { return false; - if (lhsYieldInfo->bitWidth != rhsYieldInfo->bitWidth) + } + if (lhsYieldInfo->bitWidth != rhsYieldInfo->bitWidth) { return false; + } Value lhsRecurrenceInput = lhsYieldInfo->scalar; Value rhsRecurrenceInput = rhsYieldInfo->scalar; - if (auto indexCast = lhsRecurrenceInput.getDefiningOp()) + if (auto indexCast = lhsRecurrenceInput.getDefiningOp()) { lhsRecurrenceInput = indexCast.getIn(); - if (auto indexCast = rhsRecurrenceInput.getDefiningOp()) + } + if (auto indexCast = rhsRecurrenceInput.getDefiningOp()) { rhsRecurrenceInput = indexCast.getIn(); + } // Stay conservative on unsupported cyclic proofs. The only accepted // recurrence cycle is the direct iter_arg -> plt.scalar_out self recursion // for the same value pair, optionally bridged by the index casts required by // the plt/scf type boundary; more complex cycles remain unsupported. - if (areSameValuePair(lhs, rhs, lhsRecurrenceInput, rhsRecurrenceInput)) + if (areSameValuePair(lhs, rhs, lhsRecurrenceInput, rhsRecurrenceInput)) { return true; + } return areEquivalentValues(lhsYieldInfo->scalar, rhsYieldInfo->scalar, context); @@ -207,39 +219,51 @@ static bool areEquivalentLoopCarriedValues(Value lhs, Value rhs, static bool areEquivalentOperations(Operation *lhs, Operation *rhs, ValueEquivalenceContext &context) { - if (!lhs || !rhs) + if (!lhs || !rhs) { return false; - if (lhs->getName() != rhs->getName()) + } + if (lhs->getName() != rhs->getName()) { return false; - if (lhs->getNumRegions() != 0 || rhs->getNumRegions() != 0) + } + if (lhs->getNumRegions() != 0 || rhs->getNumRegions() != 0) { return false; - if (lhs->getNumResults() != rhs->getNumResults()) + } + if (lhs->getNumResults() != rhs->getNumResults()) { return false; - if (lhs->getNumOperands() != rhs->getNumOperands()) + } + if (lhs->getNumOperands() != rhs->getNumOperands()) { return false; - if (lhs->getAttrDictionary() != rhs->getAttrDictionary()) + } + if (lhs->getAttrDictionary() != rhs->getAttrDictionary()) { return false; - if (!isMemoryEffectFree(lhs) || !isMemoryEffectFree(rhs)) + } + if (!isMemoryEffectFree(lhs) || !isMemoryEffectFree(rhs)) { return false; - if (!llvm::equal(lhs->getResultTypes(), rhs->getResultTypes())) + } + if (!llvm::equal(lhs->getResultTypes(), rhs->getResultTypes())) { return false; + } for (auto [lhsOperand, rhsOperand] : llvm::zip(lhs->getOperands(), rhs->getOperands())) { - if (!areEquivalentValues(lhsOperand, rhsOperand, context)) + if (!areEquivalentValues(lhsOperand, rhsOperand, context)) { return false; + } } return true; } static bool areEquivalentValues(Value lhs, Value rhs, ValueEquivalenceContext &context) { - if (lhs == rhs) + if (lhs == rhs) { return true; - if (!lhs || !rhs) + } + if (!lhs || !rhs) { return false; - if (lhs.getType() != rhs.getType()) + } + if (lhs.getType() != rhs.getType()) { return false; + } if (std::optional state = lookupEquivalenceState(context, lhs, rhs)) { @@ -273,13 +297,15 @@ static void populateDominatingCandidateIndices( MutableArrayRef candidates, DominanceInfo &dominanceInfo) { for (unsigned current = 0; current < candidates.size(); ++current) { for (unsigned previous = 0; previous < current; ++previous) { - if (candidates[previous].bitWidth != candidates[current].bitWidth) + if (candidates[previous].bitWidth != candidates[current].bitWidth) { continue; + } // Only reuse an earlier plt when its whole result pair dominates the // later one. This keeps replacement local and SSA-safe. if (!dominanceInfo.properlyDominates(candidates[previous].op, - candidates[current].op)) + candidates[current].op)) { continue; + } candidates[current].dominatingCandidates.push_back(previous); } } @@ -292,8 +318,9 @@ buildFusionRegionPredicateContext(pto::FusionRegionOp fusionRegion, context.fusionRegion = fusionRegion; fusionRegion.walk([&](Operation *op) -> WalkResult { - if (op != fusionRegion.getOperation() && isa(op)) + if (op != fusionRegion.getOperation() && isa(op)) { return WalkResult::skip(); + } if (std::optional candidate = buildPltCandidate(op)) context.pltCandidates.push_back(std::move(*candidate)); @@ -316,8 +343,9 @@ findEquivalentDominatingCandidate(FusionRegionPredicateContext &context, const PltCandidate ¤t = context.pltCandidates[currentIndex]; Value currentScalar = getCurrentScalarOperand(current); for (unsigned previousIndex : current.dominatingCandidates) { - if (erased.contains(previousIndex)) + if (erased.contains(previousIndex)) { continue; + } const PltCandidate &previous = context.pltCandidates[previousIndex]; // Equivalence is checked on the scalar input; when it holds, both plt // results are reused as a pair. @@ -337,14 +365,16 @@ elideEquivalentPltCandidates(FusionRegionPredicateContext &context) { for (unsigned currentIndex = 0; currentIndex < context.pltCandidates.size(); ++currentIndex) { - if (erased.contains(currentIndex)) + if (erased.contains(currentIndex)) { continue; + } std::optional previousIndex = findEquivalentDominatingCandidate(context, valueContext, currentIndex, erased); - if (!previousIndex) + if (!previousIndex) { continue; + } PltCandidate ¤t = context.pltCandidates[currentIndex]; PltCandidate &previous = context.pltCandidates[*previousIndex]; @@ -355,8 +385,9 @@ elideEquivalentPltCandidates(FusionRegionPredicateContext &context) { changed = true; } - for (Operation *op : opsToErase) + for (Operation *op : opsToErase) { op->erase(); + } return changed; } @@ -369,24 +400,28 @@ struct PTOFusionPredicateElisionPass void runOnOperation() override { func::FuncOp func = getOperation(); - if (func.isExternal()) + if (func.isExternal()) { return; + } DominanceInfo &dominanceInfo = getAnalysis(); SmallVector fusionContexts; func.walk([&](pto::FusionRegionOp fusionRegion) { FusionRegionPredicateContext context = buildFusionRegionPredicateContext(fusionRegion, dominanceInfo); - if (!context.pltCandidates.empty()) + if (!context.pltCandidates.empty()) { fusionContexts.push_back(std::move(context)); + } }); bool changed = false; - for (FusionRegionPredicateContext &context : fusionContexts) + for (FusionRegionPredicateContext &context : fusionContexts) { changed |= elideEquivalentPltCandidates(context); + } - if (!changed) + if (!changed) { markAllAnalysesPreserved(); + } } }; diff --git a/lib/PTO/Transforms/TileFusion/PTOFusionRegionGen.cpp b/lib/PTO/Transforms/TileFusion/PTOFusionRegionGen.cpp index d3de2e53f5..55f44dbf93 100644 --- a/lib/PTO/Transforms/TileFusion/PTOFusionRegionGen.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOFusionRegionGen.cpp @@ -69,8 +69,9 @@ struct PreFusionAnalysisIndex { static std::optional getRequiredI64Attr(Operation *op, StringRef attrName) { - if (auto attr = op->getAttrOfType(attrName)) + if (auto attr = op->getAttrOfType(attrName)) { return attr.getInt(); + } return std::nullopt; } @@ -87,8 +88,9 @@ collectGroupSpansInBlock(Block &block, SmallVectorImpl &spans) { GroupSpan current; auto flush = [&]() -> LogicalResult { - if (current.members.empty()) + if (current.members.empty()) { return success(); + } current.block = █ auto [it, inserted] = @@ -115,8 +117,9 @@ collectGroupSpansInBlock(Block &block, SmallVectorImpl &spans) { std::optional groupId = getRequiredI64Attr(&op, kFusionGroupIdAttr); if (!groupId) { - if (failed(flush())) + if (failed(flush())) { return failure(); + } continue; } @@ -133,8 +136,9 @@ collectGroupSpansInBlock(Block &block, SmallVectorImpl &spans) { } if (current.groupId != *groupId) { - if (failed(flush())) + if (failed(flush())) { return failure(); + } current.groupId = *groupId; } @@ -166,8 +170,9 @@ static bool isNestedInSpan(Operation *op, const DenseSet &spanOps) static void appendUniqueValue(SmallVectorImpl &values, DenseSet &seen, Value value) { - if (seen.insert(value).second) + if (seen.insert(value).second) { values.push_back(value); + } } static Operation *getTopLevelAncestorInBlock(Operation *op, Block *block) { @@ -180,8 +185,9 @@ static Operation *getTopLevelAncestorInBlock(Operation *op, Block *block) { static bool canReplaceUseWithRegionResult(OpOperand &use, Operation *boundary) { Operation *topLevel = getTopLevelAncestorInBlock(use.getOwner(), boundary->getBlock()); - if (!topLevel || topLevel == boundary) + if (!topLevel || topLevel == boundary) { return false; + } return boundary->isBeforeInBlock(topLevel); } @@ -189,10 +195,12 @@ static bool hasReplaceableUseOutsideSpan(Value value, const DenseSet &spanOps, Operation *boundary) { for (OpOperand &use : value.getUses()) { - if (isNestedInSpan(use.getOwner(), spanOps)) + if (isNestedInSpan(use.getOwner(), spanOps)) { continue; - if (canReplaceUseWithRegionResult(use, boundary)) + } + if (canReplaceUseWithRegionResult(use, boundary)) { return true; + } } return false; } @@ -207,30 +215,36 @@ static bool hasAnyUseOutsideSpan(Value value, static const FusionBlockAnalysisIndex * getBlockAnalysisIndex(const PreFusionAnalysisIndex *analysisIndex, Block *block) { - if (!analysisIndex) + if (!analysisIndex) { return nullptr; + } auto it = analysisIndex->blocks.find(block); - if (it == analysisIndex->blocks.end()) + if (it == analysisIndex->blocks.end()) { return nullptr; + } return &it->second; } static const pto::FusionWriteInstanceLiveness * getProducedWriteInstance(const FusionBlockAnalysisIndex *blockAnalysis, Operation *op, unsigned tileOutputIndex) { - if (!blockAnalysis) + if (!blockAnalysis) { return nullptr; + } auto nodeIt = blockAnalysis->nodeIdByOp.find(op); - if (nodeIt == blockAnalysis->nodeIdByOp.end()) + if (nodeIt == blockAnalysis->nodeIdByOp.end()) { return nullptr; + } auto writeIt = blockAnalysis->writeInstancesByProducerNode.find(nodeIt->second); - if (writeIt == blockAnalysis->writeInstancesByProducerNode.end()) + if (writeIt == blockAnalysis->writeInstancesByProducerNode.end()) { return nullptr; - if (tileOutputIndex >= writeIt->second.size()) + } + if (tileOutputIndex >= writeIt->second.size()) { return nullptr; + } return writeIt->second[tileOutputIndex]; } @@ -251,18 +265,22 @@ writeInstanceEscapesSpan(const pto::FusionWriteInstanceLiveness &writeInstance, static bool canSinkAllocTileDefToRegion(Value value, const GroupSpan &span, const DenseSet &spanOps) { auto alloc = dyn_cast_or_null(value.getDefiningOp()); - if (!alloc || alloc->getBlock() != span.block) + if (!alloc || alloc->getBlock() != span.block) { return false; + } Operation *firstOp = span.members.front().op; - if (!alloc->isBeforeInBlock(firstOp)) + if (!alloc->isBeforeInBlock(firstOp)) { return false; + } for (OpOperand &use : value.getUses()) { - if (isNestedInSpan(use.getOwner(), spanOps)) + if (isNestedInSpan(use.getOwner(), spanOps)) { continue; - if (!canReplaceUseWithRegionResult(use, firstOp)) + } + if (!canReplaceUseWithRegionResult(use, firstOp)) { return false; + } } return true; @@ -281,11 +299,13 @@ buildGroupSpanInterface(const GroupSpan &span, for (const GroupSpanMember &member : span.members) { spanOps.insert(member.op); - if (!blockAnalysis) + if (!blockAnalysis) { continue; + } auto nodeIt = blockAnalysis->nodeIdByOp.find(member.op); - if (nodeIt != blockAnalysis->nodeIdByOp.end()) + if (nodeIt != blockAnalysis->nodeIdByOp.end()) { spanNodeIds.insert(nodeIt->second); + } } for (const GroupSpanMember &member : span.members) { @@ -296,8 +316,9 @@ buildGroupSpanInterface(const GroupSpan &span, if (auto dpsIface = dyn_cast(member.op)) { unsigned tileOutputIndex = 0; for (Value init : dpsIface.getDpsInits()) { - if (!isa(init.getType())) + if (!isa(init.getType())) { continue; + } const pto::FusionWriteInstanceLiveness *writeInstance = getProducedWriteInstance(blockAnalysis, member.op, tileOutputIndex); @@ -305,12 +326,14 @@ buildGroupSpanInterface(const GroupSpan &span, bool escapesSpan = hasReplaceableUseOutsideSpan(init, spanOps, boundary); - if (writeInstance) + if (writeInstance) { escapesSpan = escapesSpan && writeInstanceEscapesSpan(*writeInstance, spanNodeIds); + } - if (escapesSpan) + if (escapesSpan) { appendUniqueValue(iface.externallyVisibleValues, seenOutputs, init); + } } } } @@ -321,15 +344,19 @@ buildGroupSpanInterface(const GroupSpan &span, for (const GroupSpanMember &member : span.members) { if (auto dpsIface = dyn_cast(member.op)) { for (Value init : dpsIface.getDpsInits()) { - if (!isa(init.getType())) + if (!isa(init.getType())) { continue; - if (!canSinkAllocTileDefToRegion(init, span, spanOps)) + } + if (!canSinkAllocTileDefToRegion(init, span, spanOps)) { continue; - if (hasAnyUseOutsideSpan(init, spanOps) && !visibleValues.contains(init)) + } + if (hasAnyUseOutsideSpan(init, spanOps) && !visibleValues.contains(init)) { continue; + } Operation *defOp = init.getDefiningOp(); - if (seenLocalDefs.insert(defOp).second) + if (seenLocalDefs.insert(defOp).second) { iface.localDefs.push_back(defOp); + } } } } @@ -395,23 +422,27 @@ getCommonSpanI64Attr(const GroupSpan &span, StringRef attrName) { static LogicalResult encapsulateGroupSpan(const GroupSpan &span, const PreFusionAnalysisIndex *analysisIndex) { - if (span.members.empty()) + if (span.members.empty()) { return success(); + } GroupSpanInterface iface = buildGroupSpanInterface(span, analysisIndex); FailureOr> commonRowUnroll = getCommonSpanI64Attr(span, kFusionRowUnrollAttr); - if (failed(commonRowUnroll)) + if (failed(commonRowUnroll)) { return failure(); + } FailureOr> commonColUnroll = getCommonSpanI64Attr(span, kFusionColUnrollAttr); - if (failed(commonColUnroll)) + if (failed(commonColUnroll)) { return failure(); + } SmallVector outputTypes; outputTypes.reserve(iface.externallyVisibleValues.size()); - for (Value output : iface.externallyVisibleValues) + for (Value output : iface.externallyVisibleValues) { outputTypes.push_back(output.getType()); + } Operation *firstOp = span.members.front().op; Location loc = firstOp->getLoc(); @@ -420,33 +451,39 @@ encapsulateGroupSpan(const GroupSpan &span, builder.create(loc, TypeRange(outputTypes)); fusionRegion->setAttr(kFusionGroupIdAttr, builder.getI64IntegerAttr(span.groupId)); - if (*commonRowUnroll) + if (*commonRowUnroll) { fusionRegion->setAttr(kFusionRowUnrollAttr, builder.getI64IntegerAttr(**commonRowUnroll)); - if (*commonColUnroll) + } + if (*commonColUnroll) { fusionRegion->setAttr(kFusionColUnrollAttr, builder.getI64IntegerAttr(**commonColUnroll)); + } Block *body = new Block(); fusionRegion.getBody().push_back(body); - for (Operation *localDef : iface.localDefs) + for (Operation *localDef : iface.localDefs) { localDef->moveBefore(body, body->end()); - for (const GroupSpanMember &member : span.members) + } + for (const GroupSpanMember &member : span.members) { member.op->moveBefore(body, body->end()); + } clearSpanFusionMetadata(span); SmallVector yieldValues; yieldValues.reserve(iface.externallyVisibleValues.size()); - for (Value output : iface.externallyVisibleValues) + for (Value output : iface.externallyVisibleValues) { yieldValues.push_back(output); + } OpBuilder bodyBuilder = OpBuilder::atBlockEnd(body); bodyBuilder.create(loc, ValueRange(yieldValues)); - if (failed(verify(fusionRegion.getOperation()))) + if (failed(verify(fusionRegion.getOperation()))) { return failure(); + } replaceEscapingUsesOutsideRegion(fusionRegion, iface.externallyVisibleValues); return success(); @@ -465,8 +502,9 @@ static LogicalResult processRegion(Region ®ion, return failure(); SmallVector spans; - if (failed(collectGroupSpansInBlock(block, spans))) + if (failed(collectGroupSpansInBlock(block, spans))) { return failure(); + } for (const GroupSpan &span : spans) if (failed(encapsulateGroupSpan(span, analysisIndex))) @@ -482,8 +520,9 @@ struct PTOFusionRegionGenPass void runOnOperation() override { func::FuncOp func = getOperation(); - if (func.isExternal()) + if (func.isExternal()) { return; + } // Reuse the shared pre-fusion dataflow graph cached by the analysis // manager (built once, by FusionPlan or lazily here). FusionRegionGen @@ -501,19 +540,22 @@ struct PTOFusionRegionGenPass PreFusionAnalysisIndex analysisIndex; for (const pto::FusionBlockAnalysis &blockAnalysis : analysis.blocks) { FusionBlockAnalysisIndex &index = analysisIndex.blocks[blockAnalysis.block]; - for (const pto::FusionComputeNode &node : blockAnalysis.computeNodes) + for (const pto::FusionComputeNode &node : blockAnalysis.computeNodes) { index.nodeIdByOp.try_emplace(node.op, node.id); + } for (const pto::FusionWriteInstanceLiveness &writeInstance : blockAnalysis.writeInstances) { - if (!writeInstance.producerNode) + if (!writeInstance.producerNode) { continue; + } index.writeInstancesByProducerNode[*writeInstance.producerNode] .push_back(&writeInstance); } } - if (failed(processRegion(func.getRegion(), &analysisIndex))) + if (failed(processRegion(func.getRegion(), &analysisIndex))) { signalPassFailure(); + } } }; diff --git a/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp b/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp index 42cbe80493..1bd1f7d735 100644 --- a/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp @@ -85,20 +85,27 @@ static bool isInterstageSetupOp(Operation *op) { } static bool areEquivalentOperations(Operation *lhs, Operation *rhs) { - if (!lhs || !rhs) + if (!lhs || !rhs) { return false; - if (lhs->getName() != rhs->getName()) + } + if (lhs->getName() != rhs->getName()) { return false; - if (lhs->getNumRegions() != 0 || rhs->getNumRegions() != 0) + } + if (lhs->getNumRegions() != 0 || rhs->getNumRegions() != 0) { return false; - if (lhs->getNumResults() != rhs->getNumResults()) + } + if (lhs->getNumResults() != rhs->getNumResults()) { return false; - if (lhs->getNumOperands() != rhs->getNumOperands()) + } + if (lhs->getNumOperands() != rhs->getNumOperands()) { return false; - if (lhs->getAttrDictionary() != rhs->getAttrDictionary()) + } + if (lhs->getAttrDictionary() != rhs->getAttrDictionary()) { return false; - if (!llvm::equal(lhs->getResultTypes(), rhs->getResultTypes())) + } + if (!llvm::equal(lhs->getResultTypes(), rhs->getResultTypes())) { return false; + } if (auto lhsDim = dyn_cast(lhs)) { auto rhsDim = cast(rhs); @@ -108,19 +115,23 @@ static bool areEquivalentOperations(Operation *lhs, Operation *rhs) { for (auto [lhsOperand, rhsOperand] : llvm::zip(lhs->getOperands(), rhs->getOperands())) { - if (!areEquivalentValues(lhsOperand, rhsOperand)) + if (!areEquivalentValues(lhsOperand, rhsOperand)) { return false; + } } return true; } static bool areEquivalentValues(Value lhs, Value rhs) { - if (lhs == rhs) + if (lhs == rhs) { return true; - if (!lhs || !rhs) + } + if (!lhs || !rhs) { return false; - if (lhs.getType() != rhs.getType()) + } + if (lhs.getType() != rhs.getType()) { return false; + } auto lhsArg = dyn_cast(lhs); auto rhsArg = dyn_cast(rhs); @@ -143,8 +154,9 @@ static Value traceAliasRootOneStep(Value value) { } Operation *def = value.getDefiningOp(); - if (!def) + if (!def) { return {}; + } if (auto subview = dyn_cast(def)) return subview.getSource(); @@ -163,22 +175,26 @@ static Value traceAliasRootOneStep(Value value) { if (auto reshape = dyn_cast(def)) return reshape.getSrc(); if (auto cast = dyn_cast(def)) { - if (cast.getInputs().empty()) + if (cast.getInputs().empty()) { return {}; + } if (auto result = dyn_cast(value)) { unsigned resultNumber = result.getResultNumber(); - if (resultNumber < cast.getInputs().size()) + if (resultNumber < cast.getInputs().size()) { return cast.getInputs()[resultNumber]; + } } - if (cast.getInputs().size() == 1) + if (cast.getInputs().size() == 1) { return cast.getInputs().front(); + } return {}; } if (auto forOp = dyn_cast(def)) { if (auto result = dyn_cast(value)) { unsigned resultNumber = result.getResultNumber(); - if (resultNumber < forOp.getInitArgs().size()) + if (resultNumber < forOp.getInitArgs().size()) { return forOp.getInitArgs()[resultNumber]; + } } } @@ -189,41 +205,48 @@ static Value traceAliasRoot(Value value) { int loopBound = 256; while (value) { Value upward = traceAliasRootOneStep(value); - if (!upward) + if (!upward) { break; + } value = upward; - if (loopBound-- <= 0) + if (loopBound-- <= 0) { break; + } } return value; } static LogicalResult collectAliasRelevantRoots( Operation *op, SmallVectorImpl &roots) { - if (isMemoryEffectFree(op)) + if (isMemoryEffectFree(op)) { return success(); + } auto effectsOp = dyn_cast(op); - if (!effectsOp) + if (!effectsOp) { return failure(); + } SmallVector, 4> effects; effectsOp.getEffects(effects); for (const auto &effect : effects) { Value effectValue = effect.getValue(); - if (!effectValue) + if (!effectValue) { return failure(); + } Type effectType = effectValue.getType(); if (!isa(effectType)) { - if (isa(effect.getEffect())) + if (isa(effect.getEffect())) { return failure(); + } continue; } Value root = traceAliasRoot(effectValue); - if (!root) + if (!root) { return failure(); + } roots.push_back(root); } return success(); @@ -343,16 +366,18 @@ static LogicalResult analyzeStage(scf::ForOp outerLoop, StageInfo &stage) { for (Operation &op : currentLoop.getBody()->without_terminator()) { bodyOps.push_back(&op); if (auto nestedLoop = dyn_cast(op)) { - if (childLoop) + if (childLoop) { return failure(); + } childLoop = nestedLoop; } } if (!childLoop) { for (Operation *op : bodyOps) { - if (!isSupportedLeafOp(op)) + if (!isSupportedLeafOp(op)) { return failure(); + } stage.leafOps.push_back(op); } return failure(stage.leafOps.empty()); @@ -366,16 +391,18 @@ static LogicalResult analyzeStage(scf::ForOp outerLoop, StageInfo &stage) { } if (!seenChildLoop) { // Ops before the child loop are prelude ops. - if (!isSupportedPreludeOp(op)) + if (!isSupportedPreludeOp(op)) { return failure(); + } currentLevel.preludeOps.push_back(op); } else { // Ops after the child loop are epilogue ops (e.g. row-reduction // result stores in trowmax/trowsum). They must be supported // leaf-like ops (no regions) so we can clone them into the fused // loop after all inner body ops. - if (!isSupportedPreludeOp(op)) + if (!isSupportedPreludeOp(op)) { return failure(); + } currentLevel.epilogueOps.push_back(op); } } @@ -392,9 +419,10 @@ static SmallVector collectStageRunFrom(scf::ForOp firstLoop, StageInfo firstStage; if (failed(analyzeStage(firstLoop, firstStage))) { - if (debugOS) + if (debugOS) { *debugOS << "[op-fusion] reject loop stage at " << firstLoop.getLoc() << ": stage analysis failed\n"; + } return stages; } stages.push_back(std::move(firstStage)); @@ -406,9 +434,10 @@ static SmallVector collectStageRunFrom(scf::ForOp firstLoop, nextStage.setupOps = pendingSetup; pendingSetup.clear(); if (failed(analyzeStage(nextLoop, nextStage))) { - if (debugOS) + if (debugOS) { *debugOS << "[op-fusion] stop stage run before " << nextLoop.getLoc() << ": next stage analysis failed\n"; + } break; } stages.push_back(std::move(nextStage)); @@ -416,9 +445,10 @@ static SmallVector collectStageRunFrom(scf::ForOp firstLoop, } if (!isInterstageSetupOp(op)) { - if (debugOS) + if (debugOS) { *debugOS << "[op-fusion] stop stage run at op " << op->getName() << "\n"; + } break; } pendingSetup.push_back(op); @@ -428,8 +458,9 @@ static SmallVector collectStageRunFrom(scf::ForOp firstLoop, } static bool sameLoopNestShape(const StageInfo &lhs, const StageInfo &rhs) { - if (lhs.getDepth() != rhs.getDepth()) + if (lhs.getDepth() != rhs.getDepth()) { return false; + } return llvm::all_of(llvm::zip(lhs.levels, rhs.levels), [](auto pair) { return sameForHeader(std::get<0>(pair).loop, std::get<1>(pair).loop); }); @@ -445,8 +476,9 @@ static void cloneOpAndMapResults(OpBuilder &builder, Operation *op, static void appendMappedValues(ValueRange values, IRMapping &mapping, SmallVectorImpl &mappedValues) { - for (Value value : values) + for (Value value : values) { mappedValues.push_back(mapValueOrSelf(value, mapping)); + } } static scf::ForOp buildFusedLoopNestAtLevel(OpBuilder &builder, @@ -535,22 +567,25 @@ static scf::ForOp buildFusedLoopNestAtLevel(OpBuilder &builder, static bool fuseStageRun(SmallVectorImpl &stages, llvm::raw_ostream *debugOS) { if (stages.size() < 2) { - if (debugOS) + if (debugOS) { *debugOS << "[op-fusion] reject loop run: need at least 2 stages, got " << stages.size() << "\n"; + } return false; } StageInfo &first = stages.front(); for (StageInfo &stage : llvm::drop_begin(stages)) { if (!sameLoopNestShape(first, stage)) { - if (debugOS) + if (debugOS) { *debugOS << "[op-fusion] reject loop run: loop nest shape mismatch\n"; + } return false; } } - if (!arePreludeReordersLegal(stages, debugOS)) + if (!arePreludeReordersLegal(stages, debugOS)) { return false; + } OpBuilder blockBuilder(first.getOuterLoop()); SmallVector stageMappings(stages.size()); @@ -561,8 +596,9 @@ static bool fuseStageRun(SmallVectorImpl &stages, for (Operation *setupOp : stage.setupOps) setupOp->moveBefore(fusedOuterLoop); - for (StageInfo &stage : llvm::reverse(stages)) + for (StageInfo &stage : llvm::reverse(stages)) { stage.getOuterLoop().erase(); + } return true; } @@ -575,13 +611,15 @@ static bool fuseStageRunsInBlock(Block &block, llvm::raw_ostream *debugOS) { localChange = false; for (Operation &op : block) { auto firstLoop = dyn_cast(op); - if (!firstLoop) + if (!firstLoop) { continue; + } SmallVector stages = collectStageRunFrom(firstLoop, debugOS); - if (!fuseStageRun(stages, debugOS)) + if (!fuseStageRun(stages, debugOS)) { continue; + } changed = true; localChange = true; @@ -606,19 +644,23 @@ struct PTOLowLevelLoopFusionPass int fusedFuncs = 0; for (func::FuncOp func : module.getOps()) { - if (func.isExternal()) + if (func.isExternal()) { continue; - if (func.getSymName().starts_with("__pto_oplib_")) + } + if (func.getSymName().starts_with("__pto_oplib_")) { continue; - if (func.empty()) + } + if (func.empty()) { continue; + } bool changed = false; func.walk([&](pto::FusionRegionOp fusionRegion) { changed |= fuseStageRunsInBlock(fusionRegion.getBody().front(), traceOS); }); - if (changed) + if (changed) { ++fusedFuncs; + } } if (traceEnabled) { diff --git a/lib/PTO/Transforms/TileFusion/PTOMarkLastUse.cpp b/lib/PTO/Transforms/TileFusion/PTOMarkLastUse.cpp index 32bdc83d7e..42dcb322aa 100644 --- a/lib/PTO/Transforms/TileFusion/PTOMarkLastUse.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOMarkLastUse.cpp @@ -79,16 +79,18 @@ static bool isTileInputOperand(OpOperand &operand) { static SmallVector collectTileOperands(Operation *op) { SmallVector tileOperands; for (OpOperand &operand : op->getOpOperands()) { - if (isTileOperand(operand)) + if (isTileOperand(operand)) { tileOperands.push_back(&operand); + } } return tileOperands; } static std::optional getRequiredI64Attr(Operation *op, StringRef attrName) { - if (auto attr = op->getAttrOfType(attrName)) + if (auto attr = op->getAttrOfType(attrName)) { return attr.getInt(); + } return std::nullopt; } @@ -138,8 +140,9 @@ collectGroupSpansInBlock(Block &block, SmallVectorImpl &spans) { // different group. continue; } - if (failed(flush())) + if (failed(flush())) { return failure(); + } continue; } @@ -175,17 +178,21 @@ collectGroupSpansInBlock(Block &block, SmallVectorImpl &spans) { static bool isSpanLocalLastUseCandidate(Value value, Operation *currentOp, Block *block) { - if (!value) + if (!value) { return false; + } for (OpOperand &use : value.getUses()) { Operation *user = use.getOwner(); - if (user == currentOp) + if (user == currentOp) { continue; - if (user->getBlock() != block) + } + if (user->getBlock() != block) { return false; - if (currentOp->isBeforeInBlock(user)) + } + if (currentOp->isBeforeInBlock(user)) { return false; + } } return true; } @@ -193,43 +200,51 @@ static bool isSpanLocalLastUseCandidate(Value value, Operation *currentOp, static bool hasLaterUseAfterSpan(Value value, Operation *spanEnd, Block *block) { for (OpOperand &use : value.getUses()) { Operation *user = use.getOwner(); - if (user->getBlock() != block) + if (user->getBlock() != block) { return true; - if (spanEnd->isBeforeInBlock(user)) + } + if (spanEnd->isBeforeInBlock(user)) { return true; + } } return false; } static bool isHardSpanBarrier(Operation *op) { - if (op->hasTrait() || !op->getRegions().empty()) + if (op->hasTrait() || !op->getRegions().empty()) { return true; - if (isa(op)) + } + if (isa(op)) { return true; + } return false; } static bool hasHardBarrierInSpan(const GroupSpan &span) { - if (span.members.size() < 2) + if (span.members.size() < 2) { return false; + } for (size_t i = 0; i + 1 < span.members.size(); ++i) { Operation *cur = span.members[i].op; Operation *next = span.members[i + 1].op; for (Operation *cursor = cur->getNextNode(); cursor && cursor != next; cursor = cursor->getNextNode()) { - if (isHardSpanBarrier(cursor)) + if (isHardSpanBarrier(cursor)) { return true; + } } } return false; } static void markGroupSpanLastUse(const GroupSpan &span) { - if (span.members.empty()) + if (span.members.empty()) { return; + } - if (hasHardBarrierInSpan(span)) + if (hasHardBarrierInSpan(span)) { return; + } Block &block = *span.block; Operation *spanEnd = span.members.back().op; @@ -265,15 +280,18 @@ static void markGroupSpanLastUse(const GroupSpan &span) { static LogicalResult markRegionLastUse(Region ®ion) { for (Block &block : region.getBlocks()) { SmallVector spans; - if (failed(collectGroupSpansInBlock(block, spans))) + if (failed(collectGroupSpansInBlock(block, spans))) { return failure(); - for (const GroupSpan &span : spans) + } + for (const GroupSpan &span : spans) { markGroupSpanLastUse(span); + } for (Operation &op : block) for (Region &nestedRegion : op.getRegions()) - if (failed(markRegionLastUse(nestedRegion))) + if (failed(markRegionLastUse(nestedRegion))) { return failure(); + } } return success(); } @@ -285,8 +303,9 @@ struct PTOMarkLastUsePass void runOnOperation() override { func::FuncOp func = getOperation(); - if (func.isExternal()) + if (func.isExternal()) { return; + } if (failed(markRegionLastUse(func.getRegion()))) { signalPassFailure(); diff --git a/lib/PTO/Transforms/TileFusion/PTOOpScheduling.cpp b/lib/PTO/Transforms/TileFusion/PTOOpScheduling.cpp index 941b908057..39af4e4e43 100644 --- a/lib/PTO/Transforms/TileFusion/PTOOpScheduling.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOOpScheduling.cpp @@ -59,8 +59,9 @@ struct ScheduledGroup { static std::optional getRequiredI64Attr(Operation *op, StringRef attrName) { - if (auto attr = op->getAttrOfType(attrName)) + if (auto attr = op->getAttrOfType(attrName)) { return attr.getInt(); + } return std::nullopt; } @@ -96,8 +97,9 @@ static SchedulingBarrierKind classifySchedulingBarrier(Operation *op) { return SchedulingBarrierKind::HardBoundary; } } - if (!isMemoryEffectFree(op)) + if (!isMemoryEffectFree(op)) { return SchedulingBarrierKind::HardBoundary; + } return SchedulingBarrierKind::Movable; } @@ -109,8 +111,9 @@ static bool hasTileDependency(Operation *opA, Operation *opB) { FailureOr aSemOr = pto::getFusionOpSemantics(opA); FailureOr bSemOr = pto::getFusionOpSemantics(opB); - if (failed(aSemOr) || failed(bSemOr)) + if (failed(aSemOr) || failed(bSemOr)) { return true; + } const pto::FusionOpSemantics &a = *aSemOr; const pto::FusionOpSemantics &b = *bSemOr; @@ -123,15 +126,17 @@ static bool hasTileDependency(Operation *opA, Operation *opB) { static bool crossesOperandDefinition(Operation *movingOp, Operation *candidate) { for (Value operand : movingOp->getOperands()) { Operation *defOp = operand.getDefiningOp(); - if (defOp == candidate) + if (defOp == candidate) { return true; + } } return false; } static bool canMoveEarlierAcross(Operation *movingOp, Operation *candidate) { - if (crossesOperandDefinition(movingOp, candidate)) + if (crossesOperandDefinition(movingOp, candidate)) { return false; + } switch (classifySchedulingBarrier(candidate)) { case SchedulingBarrierKind::Movable: @@ -146,8 +151,9 @@ static bool canMoveEarlierAcross(Operation *movingOp, Operation *candidate) { static bool canMoveLaterAcross(Operation *movingOp, Operation *candidate) { for (Value operand : candidate->getOperands()) { Operation *defOp = operand.getDefiningOp(); - if (defOp == movingOp) + if (defOp == movingOp) { return false; + } } switch (classifySchedulingBarrier(candidate)) { @@ -161,15 +167,18 @@ static bool canMoveLaterAcross(Operation *movingOp, Operation *candidate) { } static bool canMoveAfter(Operation *movingOp, Operation *anchorOp) { - if (!movingOp || !anchorOp || movingOp == anchorOp) + if (!movingOp || !anchorOp || movingOp == anchorOp) { return false; - if (movingOp->getBlock() != anchorOp->getBlock()) + } + if (movingOp->getBlock() != anchorOp->getBlock()) { return false; + } Operation *cursor = anchorOp->getNextNode(); while (cursor && cursor != movingOp) { - if (!canMoveEarlierAcross(movingOp, cursor)) + if (!canMoveEarlierAcross(movingOp, cursor)) { return false; + } cursor = cursor->getNextNode(); } return cursor == movingOp; @@ -214,15 +223,17 @@ collectScheduledGroups(Block &block, SmallVectorImpl &groups) { } llvm::sort(groups, [](const ScheduledGroup &lhs, const ScheduledGroup &rhs) { - if (lhs.firstOriginalIndex != rhs.firstOriginalIndex) + if (lhs.firstOriginalIndex != rhs.firstOriginalIndex) { return lhs.firstOriginalIndex < rhs.firstOriginalIndex; + } return lhs.groupId < rhs.groupId; }); for (ScheduledGroup &group : groups) { llvm::sort(group.members, [](const GroupMember &lhs, const GroupMember &rhs) { - if (lhs.order != rhs.order) + if (lhs.order != rhs.order) { return lhs.order < rhs.order; + } return lhs.originalIndex < rhs.originalIndex; }); @@ -249,10 +260,12 @@ collectScheduledGroups(Block &block, SmallVectorImpl &groups) { static bool canPrefixMoveLaterAcross( ArrayRef members, Operation *placement, Operation *barrier) { for (const GroupMember &prevMember : members) { - if (!canMoveLaterAcross(prevMember.op, barrier)) + if (!canMoveLaterAcross(prevMember.op, barrier)) { return false; - if (prevMember.op == placement) + } + if (prevMember.op == placement) { break; + } } return true; } @@ -264,14 +277,16 @@ static void movePrefixPastBarrier(ArrayRef members, for (const GroupMember &prevMember : members) { prevMember.op->moveAfter(anchor); anchor = prevMember.op; - if (prevMember.op == placement) + if (prevMember.op == placement) { break; + } } } static void scheduleGroup(ScheduledGroup &group) { - if (group.members.size() < 2) + if (group.members.size() < 2) { return; + } Operation *placement = group.members.front().op; for (GroupMember &member : llvm::drop_begin(group.members)) { @@ -287,8 +302,9 @@ static void scheduleGroup(ScheduledGroup &group) { !canMoveLaterAcross(placement, blockingOp)) break; - if (!canPrefixMoveLaterAcross(group.members, placement, blockingOp)) + if (!canPrefixMoveLaterAcross(group.members, placement, blockingOp)) { break; + } movePrefixPastBarrier(group.members, placement, blockingOp); } @@ -301,8 +317,9 @@ static LogicalResult scheduleRegion(Region ®ion) { SmallVector groups; if (failed(collectScheduledGroups(block, groups))) return failure(); - for (ScheduledGroup &group : groups) + for (ScheduledGroup &group : groups) { scheduleGroup(group); + } for (Operation &op : block) for (Region &nestedRegion : op.getRegions()) @@ -351,8 +368,9 @@ collectPhysicalFusionSpans(Block &block, bool hasCurrent = false; auto flush = [&]() { - if (!hasCurrent) + if (!hasCurrent) { return; + } spans.push_back(std::move(current)); current = FusionSpan{}; hasCurrent = false; @@ -394,8 +412,9 @@ static LogicalResult normalizeBlockFusionMetadata(Block &block, MLIRContext *context, int64_t &nextGroupId) { SmallVector spans; - if (failed(collectPhysicalFusionSpans(block, spans))) + if (failed(collectPhysicalFusionSpans(block, spans))) { return failure(); + } // Remove all old metadata before assigning canonical groups, so that // surviving span ids never collide with stale ids left on singleton or @@ -407,8 +426,9 @@ normalizeBlockFusionMetadata(Block &block, MLIRContext *context, const IntegerType i64 = IntegerType::get(context, 64); for (FusionSpan &span : spans) { - if (span.members.size() < 2) + if (span.members.size() < 2) { continue; + } const int64_t newGroupId = nextGroupId++; for (auto [order, op] : llvm::enumerate(span.members)) { @@ -425,8 +445,9 @@ static LogicalResult normalizeScheduledFusionMetadata(Region ®ion, MLIRContext *context, int64_t &nextGroupId) { for (Block &block : region.getBlocks()) { - if (failed(normalizeBlockFusionMetadata(block, context, nextGroupId))) + if (failed(normalizeBlockFusionMetadata(block, context, nextGroupId))) { return failure(); + } // Recurse into nested regions in the same pre-order walk used by // scheduleRegion, so the function-wide id counter stays deterministic. @@ -445,8 +466,9 @@ struct OpSchedulingPass void runOnOperation() override { func::FuncOp func = getOperation(); - if (func.isExternal()) + if (func.isExternal()) { return; + } if (failed(scheduleRegion(func.getRegion()))) { signalPassFailure(); diff --git a/lib/PTO/Transforms/TileFusion/PTOPreFusionAnalysis.cpp b/lib/PTO/Transforms/TileFusion/PTOPreFusionAnalysis.cpp index 1581c8297c..73b1b1b76a 100644 --- a/lib/PTO/Transforms/TileFusion/PTOPreFusionAnalysis.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOPreFusionAnalysis.cpp @@ -31,8 +31,9 @@ struct PreFusionAnalysisPass void runOnOperation() override { func::FuncOp func = getOperation(); - if (func.isExternal()) + if (func.isExternal()) { return; + } const auto &analysis = getAnalysis(); if (!analysis.isValid()) { diff --git a/lib/PTO/Transforms/TileFusion/PTOPrintPreFusionAnalysis.cpp b/lib/PTO/Transforms/TileFusion/PTOPrintPreFusionAnalysis.cpp index fcebf1c909..4cd64164b4 100644 --- a/lib/PTO/Transforms/TileFusion/PTOPrintPreFusionAnalysis.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOPrintPreFusionAnalysis.cpp @@ -87,8 +87,9 @@ static StringRef stringifyWriteInstanceEscapeClass( static void appendIndexList(llvm::raw_ostream &os, ArrayRef values) { os << "["; for (auto [idx, value] : llvm::enumerate(values)) { - if (idx) + if (idx) { os << ", "; + } os << value; } os << "]"; @@ -105,11 +106,12 @@ static void appendOptionalIndex(llvm::raw_ostream &os, static void appendDomain(llvm::raw_ostream &os, const pto::IterationDomainInfo &info) { - auto printDim = [&](int64_t dim) { - if (dim == ShapedType::kDynamic) + auto printDim = [&os](int64_t dim) { + if (dim == ShapedType::kDynamic) { os << "?"; - else + } else { os << dim; + } }; os << "("; @@ -145,8 +147,9 @@ buildValueLabels(Block &block, const pto::FusionBlockAnalysis &analysis) { for (Operation &op : block) { FailureOr semanticsOr = pto::getFusionOpSemantics(&op); - if (failed(semanticsOr)) + if (failed(semanticsOr)) { continue; + } if (semanticsOr->kind == pto::FusionOpKind::LocalBoundary) { for (Value input : semanticsOr->tileInputs) @@ -201,8 +204,9 @@ struct PrintPreFusionAnalysisPass void runOnOperation() override { func::FuncOp func = getOperation(); - if (func.isExternal()) + if (func.isExternal()) { return; + } const auto &analysis = getAnalysis(); if (!analysis.isValid()) { diff --git a/lib/PTO/Transforms/TileFusion/PTOUnrollAfterLoopFusion.cpp b/lib/PTO/Transforms/TileFusion/PTOUnrollAfterLoopFusion.cpp index f5f01faddd..d78d0f6f77 100644 --- a/lib/PTO/Transforms/TileFusion/PTOUnrollAfterLoopFusion.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOUnrollAfterLoopFusion.cpp @@ -51,8 +51,9 @@ static int64_t getEffectiveFactor(pto::FusionRegionOp region, llvm::StringRef attrName) { if (auto attr = region->getAttrOfType(attrName)) { int64_t v = attr.getInt(); - if (v > 1) + if (v > 1) { return v; + } } return 0; } @@ -74,8 +75,9 @@ static std::optional getConstantTripCount(scf::ForOp forOp) { std::optional lb = getConstantIntValue(forOp.getLowerBound()); std::optional ub = getConstantIntValue(forOp.getUpperBound()); std::optional step = getConstantIntValue(forOp.getStep()); - if (!lb || !ub || !step || *step <= 0 || *ub <= *lb) + if (!lb || !ub || !step || *step <= 0 || *ub <= *lb) { return std::nullopt; + } return (*ub - *lb + *step - 1) / *step; // ceilDiv } @@ -180,8 +182,9 @@ struct PTOUnrollAfterLoopFusion SmallVector candidates; func.walk([&](scf::ForOp forOp) { candidates.push_back(forOp); }); - for (scf::ForOp forOp : candidates) + for (scf::ForOp forOp : candidates) { (void)tryUnrollLeafForOp(forOp); + } } }; diff --git a/lib/PTO/Transforms/Utils.cpp b/lib/PTO/Transforms/Utils.cpp index cc3f577d76..ec0858e987 100644 --- a/lib/PTO/Transforms/Utils.cpp +++ b/lib/PTO/Transforms/Utils.cpp @@ -122,8 +122,9 @@ inferTFillPadLoweringKindAfterMemoryPlanning(TFillPadOp op) { std::optional inferPhysicalSectionKindFromPipe(Operation *op) { auto pipeOp = dyn_cast_or_null(op); - if (!pipeOp) + if (!pipeOp) { return std::nullopt; + } switch (pipeOp.getPipe()) { case PIPE::PIPE_M: @@ -142,8 +143,9 @@ func::ReturnOp getAssumedUniqueReturnOp(func::FuncOp funcOp) { func::ReturnOp returnOp; for (Block &b : funcOp.getBody()) { if (auto candidateOp = dyn_cast(b.getTerminator())) { - if (returnOp) + if (returnOp) { return nullptr; + } returnOp = candidateOp; } } @@ -230,8 +232,9 @@ void setBaseMemRefTypeScope(Value val, AddressSpaceAttr targetMemScope) { if (auto curMemScope = dyn_cast_if_present( dyn_cast(type).getMemorySpace())) { - if (curMemScope != targetMemScope) + if (curMemScope != targetMemScope) { llvm::report_fatal_error("memref scope mismatch while propagating PTO address space"); + } return; } @@ -325,7 +328,7 @@ std::optional> getOperationAliasInfo(Operation *op) { return std::nullopt; } -Value tracebackImpl(Value memrefVal) { +static Value tracebackImpl(Value memrefVal) { // case 1: v is the iter_arg of a scf.for if (auto arg = dyn_cast(memrefVal)) { if (auto forOp = @@ -389,13 +392,14 @@ Value tracebackImpl(Value memrefVal) { return result; } -bool isAllocLikeOp(Operation *op) { - if (!op) +static bool isAllocLikeOp(Operation *op) { + if (!op) { return false; + } return isa(op) || isa(op); } -bool isAllocLikeOp(Value val) { +static bool isAllocLikeOp(Value val) { return isAllocLikeOp(val.getDefiningOp()); } @@ -411,8 +415,9 @@ std::optional getStaticTotalSize(const ArrayRef &shapes) { } uint64_t AlignUp(uint64_t lhs, uint64_t rhs) { - if (rhs == 0) + if (rhs == 0) { return lhs; + } if (lhs % rhs != 0) { lhs += rhs - (lhs % rhs); } @@ -466,7 +471,7 @@ bool isLocalBuffer(std::optional memorySpaceAttr) { llvm_unreachable("Currently only support (UB | L1 | L0C) allocation"); } -SmallVector getOpTouchBuffer(Operation *op) { +static SmallVector getOpTouchBuffer(Operation *op) { SmallVector touchBuffer; touchBuffer.insert(touchBuffer.end(), op->getResults().begin(), op->getResults().end()); @@ -496,7 +501,7 @@ ModuleOp getTopLevelModuleOp(Operation *op) { } /// Index of yielded value where is alias of targetVal. -std::optional getYieldValueIdx(Value targetVal, ValueRange yieldedValues) { +static std::optional getYieldValueIdx(Value targetVal, ValueRange yieldedValues) { auto it = std::find(yieldedValues.begin(), yieldedValues.end(), targetVal); if (it != yieldedValues.end()) { return it - yieldedValues.begin(); @@ -506,8 +511,9 @@ std::optional getYieldValueIdx(Value targetVal, ValueRange yieldedValues) { } LoopLikeOpInterface getParentLoop(Value val) { - if (!val.getDefiningOp()) + if (!val.getDefiningOp()) { return nullptr; + } // Firstly, get parent loop LoopLikeOpInterface parentLoop = @@ -518,8 +524,9 @@ LoopLikeOpInterface getParentLoop(Value val) { // Need to determine whether val is yielded by the loop. auto yieldedValues = parentLoop.getYieldedValues(); - if (yieldedValues.empty()) + if (yieldedValues.empty()) { return parentLoop; + } auto idxLoopRes = getYieldValueIdx(val, yieldedValues); if (idxLoopRes.has_value()) { @@ -530,8 +537,9 @@ LoopLikeOpInterface getParentLoop(Value val) { // Need to determine whether val is yielded by if/else. auto parentIf = val.getDefiningOp()->getParentOfType(); - if (!parentIf || parentIf.getResults().empty()) + if (!parentIf || parentIf.getResults().empty()) { return parentLoop; + } auto thenYieldOp = parentIf.thenYield(); auto thenYieldOpers = thenYieldOp.getOperands(); diff --git a/lib/PTO/Transforms/VMILayoutAssignment.cpp b/lib/PTO/Transforms/VMILayoutAssignment.cpp index 47297c389c..0e60ed9f0b 100644 --- a/lib/PTO/Transforms/VMILayoutAssignment.cpp +++ b/lib/PTO/Transforms/VMILayoutAssignment.cpp @@ -129,7 +129,7 @@ struct LayoutSolver { unsigned addDataValue(Value value) { auto type = dyn_cast(value.getType()); if (!type) - return ~0u; + return ~0U; auto [it, inserted] = dataIds.try_emplace(value, dataNodes.size()); if (inserted) { dataNodes.push_back( @@ -144,7 +144,7 @@ struct LayoutSolver { unsigned addMaskValue(Value value) { auto type = dyn_cast(value.getType()); if (!type) - return ~0u; + return ~0U; auto [it, inserted] = maskIds.try_emplace(value, maskNodes.size()); if (inserted) maskNodes.push_back( @@ -176,7 +176,7 @@ struct LayoutSolver { LogicalResult uniteDataEquivalent(Value lhs, Value rhs, Operation *op) { unsigned lhsId = addDataValue(lhs); unsigned rhsId = addDataValue(rhs); - if (lhsId == ~0u || rhsId == ~0u) + if (lhsId == ~0U || rhsId == ~0U) return success(); unsigned lhsRoot = find(lhsId); unsigned rhsRoot = find(rhsId); @@ -207,7 +207,7 @@ struct LayoutSolver { LogicalResult uniteMask(Value lhs, Value rhs, Operation *op) { unsigned lhsId = addMaskValue(lhs); unsigned rhsId = addMaskValue(rhs); - if (lhsId == ~0u || rhsId == ~0u) + if (lhsId == ~0U || rhsId == ~0U) return success(); unsigned lhsRoot = findMask(lhsId); unsigned rhsRoot = findMask(rhsId); @@ -231,7 +231,7 @@ struct LayoutSolver { setNaturalLayout(Value value, VMILayoutAttr layout, Operation *op, DataLayoutSeedPhase phase = DataLayoutSeedPhase::Other) { unsigned id = addDataValue(value); - if (id == ~0u || !layout) + if (id == ~0U || !layout) return success(); unsigned root = find(id); VMILayoutAttr existing = dataNodes[root].naturalLayout; @@ -248,7 +248,7 @@ struct LayoutSolver { setPreferredLayout(Value value, VMILayoutAttr layout, Operation *op, DataLayoutSeedPhase phase = DataLayoutSeedPhase::Other) { unsigned id = addDataValue(value); - if (id == ~0u || !layout) + if (id == ~0U || !layout) return success(); unsigned root = find(id); VMILayoutAttr existing = dataNodes[root].preferredLayout; @@ -284,7 +284,7 @@ struct LayoutSolver { bool hasDataLayoutSeed(Value value) { unsigned id = addDataValue(value); - if (id == ~0u) + if (id == ~0U) return false; DataNode &node = dataNodes[find(id)]; return static_cast(node.naturalLayout || node.preferredLayout); @@ -481,7 +481,7 @@ struct LayoutSolver { VMILayoutAttr getDataLayout(Value value) { unsigned id = addDataValue(value); - if (id == ~0u) + if (id == ~0U) return {}; unsigned root = find(id); if (dataNodes[root].naturalLayout) diff --git a/lib/PTO/Transforms/VMILayoutSupport.cpp b/lib/PTO/Transforms/VMILayoutSupport.cpp index 4ec44e2e7f..7172cd631f 100644 --- a/lib/PTO/Transforms/VMILayoutSupport.cpp +++ b/lib/PTO/Transforms/VMILayoutSupport.cpp @@ -142,9 +142,9 @@ template static constexpr ElementCountPattern G() { static constexpr ElementCountPattern anyN() { return {{}, 0, true}; } static constexpr ElementCountPattern anyG() { return anyN(); } -static constexpr MaskGranularityPattern mb8() { return {1u << 0}; } -static constexpr MaskGranularityPattern mb16() { return {1u << 1}; } -static constexpr MaskGranularityPattern mb32() { return {1u << 2}; } +static constexpr MaskGranularityPattern mb8() { return {1U << 0}; } +static constexpr MaskGranularityPattern mb16() { return {1U << 1}; } +static constexpr MaskGranularityPattern mb32() { return {1U << 2}; } static bool matchesElementBitsPattern(ElementBitsPattern pattern, int64_t bits) { @@ -180,9 +180,9 @@ static bool matchesPhysicalChunkCountPattern( static bool matchesMaskGranularityPattern(MaskGranularityPattern pattern, StringRef granularity) { - uint8_t mask = granularity == "b8" ? 1u << 0 - : granularity == "b16" ? 1u << 1 - : granularity == "b32" ? 1u << 2 + uint8_t mask = granularity == "b8" ? 1U << 0 + : granularity == "b16" ? 1U << 1 + : granularity == "b32" ? 1U << 2 : 0; return mask != 0 && (pattern.mask & mask) != 0; } @@ -2804,11 +2804,11 @@ VMILayoutSupport::getHighPriorityGroupStoreLayoutFact( "high-priority group_store table row"); } -LogicalResult getGroupReduceAddSupportImpl(VMIVRegType sourceType, - VMIMaskType maskType, - VMIVRegType resultType, - int64_t numGroups, - std::string *reason) { +static LogicalResult getGroupReduceAddSupportImpl(VMIVRegType sourceType, + VMIMaskType maskType, + VMIVRegType resultType, + int64_t numGroups, + std::string *reason) { FailureOr fact = VMILayoutSupport().getGroupReduceLayoutFactForLayouts( sourceType, maskType, resultType, numGroups, reason); diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp index 41fb5b3285..e038e01238 100644 --- a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp +++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp @@ -1598,7 +1598,6 @@ void VMILowerUnifiedToLegacyPass::runOnOperation() { continue; } } - } std::unique_ptr mlir::pto::createVMILowerUnifiedToLegacyPass() { diff --git a/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp b/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp index 3527c61c96..db1260af15 100644 --- a/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp +++ b/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp @@ -94,7 +94,7 @@ struct MaskGranularitySolver { unsigned addMaskValue(Value value) { auto type = dyn_cast(value.getType()); if (!type) - return ~0u; + return ~0U; auto [it, inserted] = maskIds.try_emplace(value, maskNodes.size()); if (inserted) { std::string granularity; @@ -115,7 +115,7 @@ struct MaskGranularitySolver { LogicalResult uniteMask(Value lhs, Value rhs, Operation *op) { unsigned lhsId = addMaskValue(lhs); unsigned rhsId = addMaskValue(rhs); - if (lhsId == ~0u || rhsId == ~0u) + if (lhsId == ~0U || rhsId == ~0U) return success(); unsigned lhsRoot = findMask(lhsId); unsigned rhsRoot = findMask(rhsId); @@ -139,7 +139,7 @@ struct MaskGranularitySolver { LogicalResult requestMask(Value mask, StringRef granularity, Operation *op) { unsigned id = addMaskValue(mask); - if (id == ~0u) + if (id == ~0U) return success(); if (granularity.empty()) return op->emitError() << kVMIDiagLayoutContractPrefix diff --git a/lib/PTO/Transforms/VMINormalizeSignlessIntToUnsigned.cpp b/lib/PTO/Transforms/VMINormalizeSignlessIntToUnsigned.cpp index 430a1829b9..dadb6f9596 100644 --- a/lib/PTO/Transforms/VMINormalizeSignlessIntToUnsigned.cpp +++ b/lib/PTO/Transforms/VMINormalizeSignlessIntToUnsigned.cpp @@ -100,7 +100,7 @@ static void insertNormalizeCastAfter(Operation *op, OpResult result, // --------------------------------------------------------------------------- struct NormalizeSignlessPattern : public RewritePattern { - NormalizeSignlessPattern(MLIRContext *ctx) + explicit NormalizeSignlessPattern(MLIRContext *ctx) : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, ctx) {} LogicalResult matchAndRewrite(Operation *op, diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 40a3812266..52253bd2ec 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -59,8 +59,9 @@ std::optional getPointStoreDistToken(Type elementType); bool isVMIType(Type type) { return isa(type); } bool containsVMIType(Type type) { - if (isVMIType(type)) + if (isVMIType(type)) { return true; + } if (auto functionType = dyn_cast(type)) return llvm::any_of(functionType.getInputs(), @@ -68,8 +69,9 @@ bool containsVMIType(Type type) { llvm::any_of(functionType.getResults(), [](Type result) { return containsVMIType(result); }); - if (auto shapedType = dyn_cast(type)) + if (auto shapedType = dyn_cast(type)) { return containsVMIType(shapedType.getElementType()); + } return false; } @@ -94,10 +96,12 @@ struct VMISupportResult { bool isSupported() const { return supported; } LogicalResult toLogicalResult(std::string *outReason = nullptr) const { - if (supported) + if (supported) { return mlir::success(); - if (outReason) + } + if (outReason) { *outReason = reason; + } return mlir::failure(); } }; @@ -107,16 +111,19 @@ bool hasVMIType(FunctionType type) { } bool hasVMIType(Attribute attr) { - if (!attr) + if (!attr) { return false; + } if (auto typeAttr = dyn_cast(attr)) - if (containsVMIType(typeAttr.getValue())) + if (containsVMIType(typeAttr.getValue())) { return true; + } if (auto typedAttr = dyn_cast(attr)) - if (containsVMIType(typedAttr.getType())) + if (containsVMIType(typedAttr.getType())) { return true; + } if (auto arrayAttr = dyn_cast(attr)) return llvm::any_of(arrayAttr, @@ -132,17 +139,21 @@ bool hasVMIType(Attribute attr) { bool hasVMIType(Operation *op) { if (auto func = dyn_cast(op)) - if (hasVMIType(func.getFunctionType())) + if (hasVMIType(func.getFunctionType())) { return true; - if (hasVMIType(op->getOperandTypes()) || hasVMIType(op->getResultTypes())) + } + if (hasVMIType(op->getOperandTypes()) || hasVMIType(op->getResultTypes())) { return true; + } for (Region ®ion : op->getRegions()) for (Block &block : region) - if (hasVMIType(block.getArgumentTypes())) + if (hasVMIType(block.getArgumentTypes())) { return true; + } for (NamedAttribute attr : op->getAttrs()) - if (hasVMIType(attr.getValue())) + if (hasVMIType(attr.getValue())) { return true; + } return false; } @@ -155,14 +166,16 @@ StringRef getTruncFRoundModeForResult(Type resultElementType) { } StringRef getTruncFRoundMode(VMITruncFOp op, Type resultElementType) { - if (auto roundingAttr = op->getAttrOfType("rounding")) + if (auto roundingAttr = op->getAttrOfType("rounding")) { return roundingAttr.getValue(); + } return getTruncFRoundModeForResult(resultElementType); } bool isLayoutAssignedVMIType(Type type) { - if (auto vregType = dyn_cast(type)) + if (auto vregType = dyn_cast(type)) { return static_cast(vregType.getLayoutAttr()); + } if (auto maskType = dyn_cast(type)) return maskType.getLayoutAttr() && VMIMaskType::isConcreteGranularity(maskType.getGranularity()); @@ -176,41 +189,49 @@ LogicalResult verifyLayoutAssignedVMITypeTree(Operation *op, Type type) { if (auto functionType = dyn_cast(type)) { for (Type input : functionType.getInputs()) - if (failed(verifyLayoutAssignedVMITypeTree(op, input))) + if (failed(verifyLayoutAssignedVMITypeTree(op, input))) { return failure(); + } for (Type result : functionType.getResults()) - if (failed(verifyLayoutAssignedVMITypeTree(op, result))) + if (failed(verifyLayoutAssignedVMITypeTree(op, result))) { return failure(); + } } - if (auto shapedType = dyn_cast(type)) + if (auto shapedType = dyn_cast(type)) { return verifyLayoutAssignedVMITypeTree(op, shapedType.getElementType()); + } return success(); } LogicalResult verifyVMIToVPTOInputAttribute(Operation *op, Attribute attr) { - if (!attr) + if (!attr) { return success(); + } if (auto typeAttr = dyn_cast(attr)) - if (failed(verifyLayoutAssignedVMITypeTree(op, typeAttr.getValue()))) + if (failed(verifyLayoutAssignedVMITypeTree(op, typeAttr.getValue()))) { return failure(); + } if (auto typedAttr = dyn_cast(attr)) - if (failed(verifyLayoutAssignedVMITypeTree(op, typedAttr.getType()))) + if (failed(verifyLayoutAssignedVMITypeTree(op, typedAttr.getType()))) { return failure(); + } if (auto arrayAttr = dyn_cast(attr)) { for (Attribute element : arrayAttr) - if (failed(verifyVMIToVPTOInputAttribute(op, element))) + if (failed(verifyVMIToVPTOInputAttribute(op, element))) { return failure(); + } } if (auto dictAttr = dyn_cast(attr)) { for (NamedAttribute namedAttr : dictAttr) - if (failed(verifyVMIToVPTOInputAttribute(op, namedAttr.getValue()))) + if (failed(verifyVMIToVPTOInputAttribute(op, namedAttr.getValue()))) { return failure(); + } } return success(); @@ -218,35 +239,42 @@ LogicalResult verifyVMIToVPTOInputAttribute(Operation *op, Attribute attr) { LogicalResult verifyVMIToVPTOInputTypes(Operation *op) { for (Type type : op->getOperandTypes()) - if (failed(verifyLayoutAssignedVMITypeTree(op, type))) + if (failed(verifyLayoutAssignedVMITypeTree(op, type))) { return failure(); + } for (Type type : op->getResultTypes()) - if (failed(verifyLayoutAssignedVMITypeTree(op, type))) + if (failed(verifyLayoutAssignedVMITypeTree(op, type))) { return failure(); + } if (auto func = dyn_cast(op)) { FunctionType functionType = func.getFunctionType(); for (Type type : functionType.getInputs()) - if (failed(verifyLayoutAssignedVMITypeTree(op, type))) + if (failed(verifyLayoutAssignedVMITypeTree(op, type))) { return failure(); + } for (Type type : functionType.getResults()) - if (failed(verifyLayoutAssignedVMITypeTree(op, type))) + if (failed(verifyLayoutAssignedVMITypeTree(op, type))) { return failure(); + } } for (Region ®ion : op->getRegions()) for (Block &block : region) for (Type type : block.getArgumentTypes()) - if (failed(verifyLayoutAssignedVMITypeTree(op, type))) + if (failed(verifyLayoutAssignedVMITypeTree(op, type))) { return failure(); + } for (NamedAttribute attr : op->getAttrs()) - if (failed(verifyVMIToVPTOInputAttribute(op, attr.getValue()))) + if (failed(verifyVMIToVPTOInputAttribute(op, attr.getValue()))) { return failure(); + } return success(); } LogicalResult verifyVMIToVPTOInputIR(ModuleOp module) { WalkResult result = module.walk([&](Operation *op) { - if (failed(verifyVMIToVPTOInputTypes(op))) + if (failed(verifyVMIToVPTOInputTypes(op))) { return WalkResult::interrupt(); + } return WalkResult::advance(); }); return failure(result.wasInterrupted()); @@ -254,8 +282,9 @@ LogicalResult verifyVMIToVPTOInputIR(ModuleOp module) { static Value materializeVPTOToVMI(OpBuilder &builder, Type resultType, ValueRange inputs, Location loc) { - if (!isVMIType(resultType)) + if (!isVMIType(resultType)) { return {}; + } return builder.create(loc, resultType, inputs).getResult(); } @@ -263,19 +292,23 @@ static SmallVector materializeVMIToVPTO(OpBuilder &builder, TypeRange resultTypes, ValueRange inputs, Location loc) { - if (inputs.size() != 1 || !isVMIType(inputs.front().getType())) + if (inputs.size() != 1 || !isVMIType(inputs.front().getType())) { return {}; + } auto unpackOp = builder.create(loc, resultTypes, inputs.front()); return SmallVector(unpackOp->getResults()); } static int64_t getMaskGranularityBits(StringRef granularity) { - if (granularity == "b8") + if (granularity == "b8") { return 8; - if (granularity == "b16") + } + if (granularity == "b16") { return 16; - if (granularity == "b32") + } + if (granularity == "b32") { return 32; + } return 0; } @@ -294,16 +327,18 @@ static StringRef getMaskGranularityForBits(int64_t bits) { static FailureOr getVMIMaskPhysicalGranularity(VMIMaskType type) { int64_t bits = getMaskGranularityBits(type.getGranularity()); - if (bits == 0) + if (bits == 0) { return failure(); + } VMILayoutAttr layout = type.getLayoutAttr(); int64_t laneStride = layout && layout.hasLaneStride() ? layout.getLaneStride() : 1; int64_t physicalBits = bits * laneStride; StringRef physicalGranularity = getMaskGranularityForBits(physicalBits); - if (physicalGranularity.empty()) + if (physicalGranularity.empty()) { return failure(); + } return physicalGranularity; } @@ -315,12 +350,14 @@ class VMIToVPTOTypeConverter final : public TypeConverter { SmallVectorImpl &results) -> LogicalResult { FailureOr arity = getVMIPhysicalArity(type); Type physicalElementType = getVMIPhysicalDataElementType(type); - if (failed(arity)) + if (failed(arity)) { return failure(); + } FailureOr lanesPerPart = getDataLanesPerPart(physicalElementType); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return failure(); + } for (int64_t i = 0; i < *arity; ++i) results.push_back(VRegType::get(type.getContext(), *lanesPerPart, physicalElementType)); @@ -331,8 +368,9 @@ class VMIToVPTOTypeConverter final : public TypeConverter { FailureOr arity = getVMIPhysicalArity(type); FailureOr physicalGranularity = getVMIMaskPhysicalGranularity(type); - if (failed(arity) || failed(physicalGranularity)) + if (failed(arity) || failed(physicalGranularity)) { return failure(); + } for (int64_t i = 0; i < *arity; ++i) results.push_back( MaskType::get(type.getContext(), *physicalGranularity)); @@ -346,8 +384,9 @@ class VMIToVPTOTypeConverter final : public TypeConverter { FailureOr> getConvertedResultTypes(Operation *op, unsigned resultIndex, const TypeConverter &typeConverter) { - if (resultIndex >= op->getNumResults()) + if (resultIndex >= op->getNumResults()) { return failure(); + } SmallVector resultTypes; if (failed(typeConverter.convertType(op->getResult(resultIndex).getType(), resultTypes))) @@ -358,8 +397,9 @@ getConvertedResultTypes(Operation *op, unsigned resultIndex, FailureOr> getConvertedResultTypes(Operation *op, const TypeConverter &typeConverter) { SmallVector resultTypes; - if (failed(typeConverter.convertTypes(op->getResultTypes(), resultTypes))) + if (failed(typeConverter.convertTypes(op->getResultTypes(), resultTypes))) { return failure(); + } return resultTypes; } @@ -369,8 +409,9 @@ getConvertedVRegTypesWithLayout(VMIVRegType type, VMILayoutAttr layout, auto relayoutType = VMIVRegType::get(type.getContext(), type.getElementCount(), type.getElementType(), layout); SmallVector convertedTypes; - if (failed(typeConverter.convertType(relayoutType, convertedTypes))) + if (failed(typeConverter.convertType(relayoutType, convertedTypes))) { return failure(); + } return convertedTypes; } @@ -378,15 +419,18 @@ FailureOr getVRegPhysicalFootprintBytes(TypeRange types) { int64_t totalBytes = 0; for (Type type : types) { auto vregType = dyn_cast(type); - if (!vregType) + if (!vregType) { return failure(); + } unsigned elementBits = pto::getPTOStorageElemBitWidth(vregType.getElementType()); - if (elementBits == 0) + if (elementBits == 0) { return failure(); + } int64_t chunkBits = vregType.getElementCount() * elementBits; - if (chunkBits % 8 != 0) + if (chunkBits % 8 != 0) { return failure(); + } totalBytes += chunkBits / 8; } return totalBytes; @@ -398,8 +442,9 @@ FailureOr hasNoWiderFootprintThanContiguous(TypeRange assignedTypes, getVRegPhysicalFootprintBytes(assignedTypes); FailureOr contiguousBytes = getVRegPhysicalFootprintBytes(contiguousTypes); - if (failed(assignedBytes) || failed(contiguousBytes)) + if (failed(assignedBytes) || failed(contiguousBytes)) { return failure(); + } return *assignedBytes <= *contiguousBytes; } @@ -431,19 +476,22 @@ void replaceOpWithFlatConvertedValues( SmallVector flattenOneToNOperands(ArrayRef operands) { SmallVector flat; - for (ValueRange operand : operands) + for (ValueRange operand : operands) { llvm::append_range(flat, operand); + } return flat; } bool isIdentityOneToNValueMapping(ValueRange originalValues, ArrayRef convertedValues) { - if (originalValues.size() != convertedValues.size()) + if (originalValues.size() != convertedValues.size()) { return false; + } for (auto [original, converted] : llvm::zip_equal(originalValues, convertedValues)) { - if (converted.size() != 1 || converted.front() != original) + if (converted.size() != 1 || converted.front() != original) { return false; + } } return true; } @@ -452,8 +500,9 @@ TypeRange getConvertedSignatureTypes( const TypeConverter::SignatureConversion &conversion, unsigned originalIndex) { TypeRange convertedTypes = conversion.getConvertedTypes(); - if (auto mapping = conversion.getInputMapping(originalIndex)) + if (auto mapping = conversion.getInputMapping(originalIndex)) { return convertedTypes.slice(mapping->inputNo, mapping->size); + } return {}; } @@ -462,8 +511,9 @@ bool hasNonIdentitySignatureConversion( const TypeConverter::SignatureConversion &conversion) { for (auto [index, originalType] : llvm::enumerate(originalTypes)) { TypeRange convertedTypes = getConvertedSignatureTypes(conversion, index); - if (convertedTypes.size() != 1 || convertedTypes.front() != originalType) + if (convertedTypes.size() != 1 || convertedTypes.front() != originalType) { return true; + } } return false; } @@ -494,12 +544,15 @@ FailureOr createAllTrueMaskForVReg(Location loc, VRegType vregType, FailureOr getMaskTypeForVReg(VRegType vregType, MLIRContext *ctx) { unsigned elementBits = pto::getPTOStorageElemBitWidth(vregType.getElementType()); - if (elementBits == 8) + if (elementBits == 8) { return MaskType::get(ctx, "b8"); - if (elementBits == 16) + } + if (elementBits == 16) { return MaskType::get(ctx, "b16"); - if (elementBits == 32) + } + if (elementBits == 32) { return MaskType::get(ctx, "b32"); + } return failure(); } @@ -563,15 +616,18 @@ FailureOr createPrefixMask(Location loc, MaskType maskType, } bool areEquivalentReductionMasks(Value lhs, Value rhs) { - if (lhs == rhs) + if (lhs == rhs) { return true; - if (lhs.getType() != rhs.getType()) + } + if (lhs.getType() != rhs.getType()) { return false; + } Operation *lhsOp = lhs.getDefiningOp(); Operation *rhsOp = rhs.getDefiningOp(); - if (!lhsOp || !rhsOp || lhsOp->getName() != rhsOp->getName()) + if (!lhsOp || !rhsOp || lhsOp->getName() != rhsOp->getName()) { return false; + } bool isPatternMask = isa( @@ -630,15 +686,17 @@ createRuntimePrefixMask(Location loc, MaskType maskType, Value activeLanes, LogicalResult checkSupportedMaskableVReg(VMIVRegType type, std::string *reason = nullptr) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; FailureOr lanesPerPart = getDataLanesPerPart(type.getElementType()); FailureOr arity = getVMIPhysicalArity(type); - if (failed(lanesPerPart) || failed(arity) || *arity < 1) + if (failed(lanesPerPart) || failed(arity) || *arity < 1) { return fail("requires computable non-empty physical vreg parts"); + } return success(); } @@ -656,8 +714,9 @@ Value createI16Constant(Location loc, int64_t value, FailureOr createPrefixMaskForActiveLanes(Location loc, MaskType maskType, int64_t activeLanes, PatternRewriter &rewriter) { - if (activeLanes <= 0) + if (activeLanes <= 0) { return createPrefixMask(loc, maskType, "PAT_ALLF", rewriter); + } switch (activeLanes) { case 1: @@ -674,8 +733,9 @@ FailureOr createPrefixMaskForActiveLanes(Location loc, MaskType maskType, default: { FailureOr> dynamicMask = createRuntimePrefixMask( loc, maskType, createI32Constant(loc, activeLanes, rewriter), rewriter); - if (failed(dynamicMask)) + if (failed(dynamicMask)) { return failure(); + } return dynamicMask->first; } } @@ -695,8 +755,9 @@ Value clampDynamicActiveLanes(Location loc, Value activeLanes, Value createPartitionActiveLanes(Location loc, Value activeLanesI32, int64_t factor, int64_t part, PatternRewriter &rewriter) { - if (factor == 1) + if (factor == 1) { return activeLanesI32; + } int64_t bias = factor - 1 - part; Value biased = activeLanesI32; if (bias != 0) @@ -707,8 +768,9 @@ Value createPartitionActiveLanes(Location loc, Value activeLanesI32, } std::optional getPowerOfTwoLog2(int64_t value) { - if (value <= 0 || (value & (value - 1)) != 0) + if (value <= 0 || (value & (value - 1)) != 0) { return std::nullopt; + } int64_t log2 = 0; while (value > 1) { value >>= 1; @@ -719,10 +781,12 @@ std::optional getPowerOfTwoLog2(int64_t value) { std::optional getPrefixPattern(int64_t activeLanes, int64_t lanesPerPart) { - if (activeLanes <= 0) + if (activeLanes <= 0) { return std::string("PAT_ALLF"); - if (activeLanes >= lanesPerPart) + } + if (activeLanes >= lanesPerPart) { return std::string("PAT_ALL"); + } switch (activeLanes) { case 1: case 2: @@ -755,16 +819,18 @@ static int64_t ceilDivNonNegative(int64_t lhs, int64_t rhs) { FailureOr getDataLayoutFactor(VMIVRegType type) { VMILayoutAttr layout = type.getLayoutAttr(); - if (!layout) + if (!layout) { return failure(); + } return layout.isDenseSplit() ? layout.getFactor() : 1; } FailureOr getDataChunksInPart(VMIVRegType type, int64_t part) { FailureOr factor = getDataLayoutFactor(type); FailureOr lanesPerPart = getDataLanesPerPart(type.getElementType()); - if (failed(factor) || failed(lanesPerPart) || part < 0 || part >= *factor) + if (failed(factor) || failed(lanesPerPart) || part < 0 || part >= *factor) { return failure(); + } int64_t logicalLanesInPart = (type.getElementCount() + *factor - 1 - part) / *factor; @@ -774,50 +840,59 @@ FailureOr getDataChunksInPart(VMIVRegType type, int64_t part) { FailureOr getDataFlatPartIndex(VMIVRegType type, int64_t part, int64_t chunk) { FailureOr factor = getDataLayoutFactor(type); - if (failed(factor) || part < 0 || part >= *factor || chunk < 0) + if (failed(factor) || part < 0 || part >= *factor || chunk < 0) { return failure(); + } int64_t flatIndex = 0; for (int64_t currentPart = 0; currentPart < part; ++currentPart) { FailureOr chunks = getDataChunksInPart(type, currentPart); - if (failed(chunks)) + if (failed(chunks)) { return failure(); + } flatIndex += *chunks; } FailureOr chunks = getDataChunksInPart(type, part); - if (failed(chunks) || chunk >= *chunks) + if (failed(chunks) || chunk >= *chunks) { return failure(); + } return flatIndex + chunk; } FailureOr checkFullDataPhysicalChunks(VMIVRegType type, std::string *reason) { auto fail = [&](const Twine &message) -> FailureOr { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; FailureOr lanesPerPart = getDataLanesPerPart(type.getElementType()); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return fail("requires known physical lanes per part"); + } FailureOr factor = getDataLayoutFactor(type); - if (failed(factor)) + if (failed(factor)) { return fail("requires assigned layout"); + } for (int64_t part = 0; part < *factor; ++part) { FailureOr chunks = getDataChunksInPart(type, part); - if (failed(chunks)) + if (failed(chunks)) { return fail("requires known physical chunks"); + } for (int64_t chunk = 0; chunk < *chunks; ++chunk) { for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { FailureOr padding = isPaddingLane(type, part, chunk, lane); - if (failed(padding)) + if (failed(padding)) { return fail("failed to map physical padding lane"); - if (*padding) + } + if (*padding) { return fail("found padding lane in physical chunk"); + } } } } @@ -827,24 +902,29 @@ FailureOr checkFullDataPhysicalChunks(VMIVRegType type, FailureOr getVMITypeLayoutFactor(Type type) { Attribute layout; - if (auto vregType = dyn_cast(type)) + if (auto vregType = dyn_cast(type)) { layout = vregType.getLayout(); + } else if (auto maskType = dyn_cast(type)) layout = maskType.getLayout(); - else + else { return failure(); + } auto layoutAttr = dyn_cast_or_null(layout); - if (!layoutAttr) + if (!layoutAttr) { return failure(); + } return layoutAttr.isDenseSplit() ? layoutAttr.getFactor() : 1; } FailureOr getVMITypeElementCount(Type type) { - if (auto vregType = dyn_cast(type)) + if (auto vregType = dyn_cast(type)) { return vregType.getElementCount(); - if (auto maskType = dyn_cast(type)) + } + if (auto maskType = dyn_cast(type)) { return maskType.getElementCount(); + } return failure(); } @@ -855,8 +935,9 @@ FailureOr getVMITypeLanesPerPart(Type type) { if (auto maskType = dyn_cast(type)) { FailureOr physicalGranularity = getVMIMaskPhysicalGranularity(maskType); - if (failed(physicalGranularity)) + if (failed(physicalGranularity)) { return failure(); + } return getMaskLanesPerPart(*physicalGranularity); } return failure(); @@ -871,17 +952,20 @@ FailureOr getVMITypeChunksInPart(Type type, int64_t part) { return failure(); VMILayoutAttr layout; - if (auto vregType = dyn_cast(type)) + if (auto vregType = dyn_cast(type)) { layout = vregType.getLayoutAttr(); + } else if (auto maskType = dyn_cast(type)) layout = maskType.getLayoutAttr(); - if (!layout) + if (!layout) { return failure(); + } int64_t logicalLanesInPart = (*elementCount + *factor - 1 - part) / *factor; int64_t laneStride = 1; - if (isa(type) && layout.isDense()) + if (isa(type) && layout.isDense()) { laneStride = layout.getLaneStride(); + } int64_t physicalLanes = logicalLanesInPart == 0 ? 0 : (logicalLanesInPart - 1) * laneStride + 1; return ceilDivNonNegative(physicalLanes, *lanesPerPart); @@ -889,27 +973,32 @@ FailureOr getVMITypeChunksInPart(Type type, int64_t part) { LogicalResult checkFullVMIPhysicalChunks(Type type, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; FailureOr factor = getVMITypeLayoutFactor(type); FailureOr lanesPerPart = getVMITypeLanesPerPart(type); - if (failed(factor) || failed(lanesPerPart)) + if (failed(factor) || failed(lanesPerPart)) { return fail("requires assigned layout with known physical lanes per part"); + } for (int64_t part = 0; part < *factor; ++part) { FailureOr chunks = getVMITypeChunksInPart(type, part); - if (failed(chunks)) + if (failed(chunks)) { return fail("requires known physical chunks"); + } for (int64_t chunk = 0; chunk < *chunks; ++chunk) { for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { FailureOr padding = isPaddingLane(type, part, chunk, lane); - if (failed(padding)) + if (failed(padding)) { return fail("failed to map physical padding lane"); - if (*padding) + } + if (*padding) { return fail("found padding lane in physical chunk"); + } } } } @@ -923,44 +1012,53 @@ FailureOr getContiguousMaterializationPartCount(Type type, FailureOr getContiguousMaterializationPartCount(Type type, std::string *reason) { auto fail = [&](const Twine &message) -> FailureOr { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; FailureOr arity = getVMIPhysicalArity(type); FailureOr factor = getVMITypeLayoutFactor(type); - if (failed(arity) || failed(factor)) + if (failed(arity) || failed(factor)) { return fail("requires computable physical arity and assigned layout"); + } Attribute layoutAttr; - if (auto vregType = dyn_cast(type)) + if (auto vregType = dyn_cast(type)) { layoutAttr = vregType.getLayout(); + } else if (auto maskType = dyn_cast(type)) layoutAttr = maskType.getLayout(); - else + else { return fail("requires VMI data or mask type"); + } auto layout = dyn_cast_or_null(layoutAttr); - if (!layout) + if (!layout) { return fail("requires assigned layout"); - if (layout.isContiguous() && layout.getLaneStride() == 1) + } + if (layout.isContiguous() && layout.getLaneStride() == 1) { return *arity; + } if (!layout.isDenseSplit() || (layout.getFactor() != 2 && layout.getFactor() != 4)) return fail("requires contiguous, deinterleaved=2/4, or " "block_deinterleaved=2/4 layout"); FailureOr chunksPerGroup = getVMITypeChunksInPart(type, 0); - if (failed(chunksPerGroup)) + if (failed(chunksPerGroup)) { return fail("requires known physical chunks per part"); - if (*chunksPerGroup == 0) + } + if (*chunksPerGroup == 0) { return fail("requires at least one physical chunk per part"); + } for (int64_t part = 1; part < *factor; ++part) { FailureOr chunks = getVMITypeChunksInPart(type, part); - if (failed(chunks)) + if (failed(chunks)) { return fail("requires known physical chunks per part"); + } if (layout.getFactor() == 2 && *chunks != *chunksPerGroup) return fail("requires every deinterleaved part to have the same " "physical chunk count"); @@ -988,11 +1086,13 @@ LogicalResult checkCanMaterializeToContiguous(Type type, std::string *reason) { } std::optional getConstantIndexValue(Value value) { - if (auto constant = value.getDefiningOp()) + if (auto constant = value.getDefiningOp()) { return constant.value(); + } if (auto constant = value.getDefiningOp()) { - if (auto integerAttr = dyn_cast(constant.getValue())) + if (auto integerAttr = dyn_cast(constant.getValue())) { return integerAttr.getInt(); + } } return std::nullopt; } @@ -1004,22 +1104,29 @@ static int64_t normalizeRemainder(int64_t value, int64_t modulus) { std::optional getKnownIndexRemainder(Value value, int64_t modulus, int depth = 0) { - if (modulus <= 1) + if (modulus <= 1) { return 0; - if (depth > 6) + } + if (depth > 6) { return std::nullopt; - if (std::optional constant = getConstantIndexValue(value)) + } + if (std::optional constant = getConstantIndexValue(value)) { return normalizeRemainder(*constant, modulus); + } - if (auto cast = value.getDefiningOp()) + if (auto cast = value.getDefiningOp()) { return getKnownIndexRemainder(cast.getIn(), modulus, depth + 1); - if (auto cast = value.getDefiningOp()) + } + if (auto cast = value.getDefiningOp()) { return getKnownIndexRemainder(cast.getIn(), modulus, depth + 1); - if (auto cast = value.getDefiningOp()) + } + if (auto cast = value.getDefiningOp()) { return getKnownIndexRemainder(cast.getIn(), modulus, depth + 1); + } if (auto cast = value.getDefiningOp()) { - if (cast->getNumOperands() == 1 && cast->getNumResults() == 1) + if (cast->getNumOperands() == 1 && cast->getNumResults() == 1) { return getKnownIndexRemainder(cast.getOperand(0), modulus, depth + 1); + } } if (auto add = value.getDefiningOp()) { @@ -1027,8 +1134,9 @@ std::optional getKnownIndexRemainder(Value value, int64_t modulus, getKnownIndexRemainder(add.getLhs(), modulus, depth + 1); std::optional rhs = getKnownIndexRemainder(add.getRhs(), modulus, depth + 1); - if (lhs && rhs) + if (lhs && rhs) { return normalizeRemainder(*lhs + *rhs, modulus); + } return std::nullopt; } if (auto sub = value.getDefiningOp()) { @@ -1036,8 +1144,9 @@ std::optional getKnownIndexRemainder(Value value, int64_t modulus, getKnownIndexRemainder(sub.getLhs(), modulus, depth + 1); std::optional rhs = getKnownIndexRemainder(sub.getRhs(), modulus, depth + 1); - if (lhs && rhs) + if (lhs && rhs) { return normalizeRemainder(*lhs - *rhs, modulus); + } return std::nullopt; } if (auto mul = value.getDefiningOp()) { @@ -1045,10 +1154,12 @@ std::optional getKnownIndexRemainder(Value value, int64_t modulus, getKnownIndexRemainder(mul.getLhs(), modulus, depth + 1); std::optional rhs = getKnownIndexRemainder(mul.getRhs(), modulus, depth + 1); - if ((lhs && *lhs == 0) || (rhs && *rhs == 0)) + if ((lhs && *lhs == 0) || (rhs && *rhs == 0)) { return 0; - if (lhs && rhs) + } + if (lhs && rhs) { return normalizeRemainder(*lhs * *rhs, modulus); + } return std::nullopt; } @@ -1058,23 +1169,28 @@ std::optional getKnownIndexRemainder(Value value, int64_t modulus, std::optional getKnownPointerByteRemainder(Value pointer, int64_t alignmentBytes, int depth = 0) { - if (alignmentBytes <= 1) + if (alignmentBytes <= 1) { return 0; - if (depth > 6) + } + if (depth > 6) { return std::nullopt; + } // A raw PTO pointer block argument is a base address. Its required // address-space alignment is an ABI precondition; derived pointers must // prove that their element offsets preserve that alignment. - if (isa(pointer)) + if (isa(pointer)) { return 0; + } if (auto cast = pointer.getDefiningOp()) { Value input = cast.getInput(); - if (isa(input.getType())) + if (isa(input.getType())) { return getKnownPointerByteRemainder(input, alignmentBytes, depth + 1); - if (isa(input.getType())) + } + if (isa(input.getType())) { return getKnownIndexRemainder(input, alignmentBytes, depth + 1); + } return std::nullopt; } @@ -1089,19 +1205,22 @@ std::optional getKnownPointerByteRemainder(Value pointer, std::optional base = getKnownPointerByteRemainder(add.getPtr(), alignmentBytes, depth + 1); auto pointerType = dyn_cast(add.getPtr().getType()); - if (!base || !pointerType) + if (!base || !pointerType) { return std::nullopt; + } unsigned elementBits = pto::getPTOStorageElemBitWidth(pointerType.getElementType()); - if (elementBits == 0 || elementBits % 8 != 0) + if (elementBits == 0 || elementBits % 8 != 0) { return std::nullopt; + } int64_t elementBytes = elementBits / 8; int64_t offsetModulus = alignmentBytes / std::gcd(alignmentBytes, elementBytes); std::optional offset = getKnownIndexRemainder(add.getOffset(), offsetModulus, depth + 1); - if (!offset) + if (!offset) { return std::nullopt; + } return normalizeRemainder(*base + *offset * elementBytes, alignmentBytes); } @@ -1114,36 +1233,42 @@ bool isKnown32ByteAlignedAddress(Value pointer, Value elementOffset, std::optional base = getKnownPointerByteRemainder(pointer, alignmentBytes); unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); - if (!base || elementBits == 0 || elementBits % 8 != 0) + if (!base || elementBits == 0 || elementBits % 8 != 0) { return false; + } int64_t elementBytes = elementBits / 8; int64_t offsetModulus = alignmentBytes / std::gcd(alignmentBytes, elementBytes); std::optional offset = getKnownIndexRemainder(elementOffset, offsetModulus); - if (!offset) + if (!offset) { return false; + } return normalizeRemainder(*base + *offset * elementBytes, alignmentBytes) == 0; } FailureOr getStaticMemRefElementCount(Type type) { auto memrefType = dyn_cast(type); - if (!memrefType || !memrefType.hasStaticShape()) + if (!memrefType || !memrefType.hasStaticShape()) { return failure(); + } int64_t elements = 1; - for (int64_t dim : memrefType.getShape()) + for (int64_t dim : memrefType.getShape()) { elements *= dim; + } return elements; } static Type getMemoryElementType(Type type) { - if (auto ptrType = dyn_cast(type)) + if (auto ptrType = dyn_cast(type)) { return ptrType.getElementType(); - if (auto memrefType = dyn_cast(type)) + } + if (auto memrefType = dyn_cast(type)) { return memrefType.getElementType(); + } return {}; } @@ -1235,16 +1360,18 @@ buildContiguousIdentityLaneAddressMap(int64_t constantOffset, VMIVRegType resultType, std::string *reason = nullptr) { auto fail = [&](const Twine &message) -> FailureOr { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; FailureOr lanesPerPart = getDataLanesPerPart(resultType.getElementType()); FailureOr arity = getVMIPhysicalArity(resultType); - if (failed(lanesPerPart) || failed(arity)) + if (failed(lanesPerPart) || failed(arity)) { return fail("requires computable physical read footprint"); + } VMIMemoryLaneAddressMap map; map.baseElementOffset = constantOffset; @@ -1255,8 +1382,9 @@ buildContiguousIdentityLaneAddressMap(int64_t constantOffset, VMISupportResult requireIdentityMemRefLayout(Type memoryType, StringRef role, Value memoryValue = {}) { auto memrefType = dyn_cast(memoryType); - if (!memrefType || memrefType.getLayout().isIdentity()) + if (!memrefType || memrefType.getLayout().isIdentity()) { return VMISupportResult::success(); + } std::string reason = (Twine(role) + " memref layout is non-identity; current VMI memory access plan " @@ -1280,24 +1408,28 @@ computeSafeFullReadProof(Type sourceType, std::optional constantOffset, return proof; }; - if (!constantOffset) + if (!constantOffset) { return fail("requires constant index offset"); + } FailureOr staticElements = getStaticMemRefElementCount(sourceType); - if (failed(staticElements)) + if (failed(staticElements)) { return fail("requires statically shaped memref source"); + } int64_t elements = *staticElements; proof.staticElementCount = elements; - if (*constantOffset < 0) + if (*constantOffset < 0) { return fail("requires non-negative offset"); + } std::string addressMapReason; FailureOr addressMap = buildContiguousIdentityLaneAddressMap(*constantOffset, resultType, &addressMapReason); - if (failed(addressMap)) + if (failed(addressMap)) { return fail(addressMapReason); + } proof.laneAddressMap = *addressMap; proof.physicalFootprint = addressMap->physicalLaneFootprint; @@ -1370,20 +1502,23 @@ FailureOr verifyFullOrSafeReadVRegChunks(Operation *op, std::string fullChunkReason; FailureOr lanesPerPart = checkFullDataPhysicalChunks(type, &fullChunkReason); - if (succeeded(lanesPerPart)) + if (succeeded(lanesPerPart)) { return *lanesPerPart; + } VMIMemorySafeReadProof safeReadProof = computeSafeFullReadProof(sourceType, getConstantIndexValue(offset), type); if (safeReadProof.proven) { lanesPerPart = getDataLanesPerPart(type.getElementType()); - if (succeeded(lanesPerPart)) + if (succeeded(lanesPerPart)) { return *lanesPerPart; + } } lanesPerPart = getDataLanesPerPart(type.getElementType()); - if (succeeded(lanesPerPart)) + if (succeeded(lanesPerPart)) { return *lanesPerPart; + } (void)rewriter.notifyMatchFailure( op, Twine("memory lowering ") + fullChunkReason + @@ -1396,26 +1531,31 @@ checkSupportedLoadShape(VMIVRegType type, Value source, Type sourceType, std::optional constantOffset, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(source, sourceType, type, constantOffset, VMIMemoryValidMaskKind::AllTrue); - if (!accessPlan.layoutSupport.isSupported()) + if (!accessPlan.layoutSupport.isSupported()) { return fail(accessPlan.layoutSupport.reason); + } VMILayoutSupport supports; - if (failed(supports.getLoadLayoutFact(type, reason))) + if (failed(supports.getLoadLayoutFact(type, reason))) { return failure(); + } - if (getDenseLaneStrideLoadDistToken(type)) + if (getDenseLaneStrideLoadDistToken(type)) { return success(); + } - if (failed(getDataLanesPerPart(type.getElementType()))) + if (failed(getDataLanesPerPart(type.getElementType()))) { return fail("requires element type with known physical lane width"); + } return success(); } @@ -1423,8 +1563,9 @@ LogicalResult checkSupportedDeinterleaveLoadShape( VMIDeinterleaveLoadOp op, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -1434,17 +1575,20 @@ LogicalResult checkSupportedDeinterleaveLoadShape( if (failed(supports.getDeinterleaveLoadLayoutFactForLayouts( lowType, highType, reason))) return failure(); - if (!getX2MemoryDistToken(lowType.getElementType(), "DINTLV")) + if (!getX2MemoryDistToken(lowType.getElementType(), "DINTLV")) { return fail("requires 8/16/32-bit element type for vldsx2 DINTLV"); + } VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(op.getSource(), op.getSource().getType(), lowType, getConstantIndexValue(op.getOffset()), VMIMemoryValidMaskKind::AllTrue); - if (!accessPlan.layoutSupport.isSupported()) + if (!accessPlan.layoutSupport.isSupported()) { return fail(accessPlan.layoutSupport.reason); + } std::string fullChunkReason; - if (failed(checkFullDataPhysicalChunks(lowType, &fullChunkReason))) + if (failed(checkFullDataPhysicalChunks(lowType, &fullChunkReason))) { return fail(Twine("requires full physical chunks; ") + fullChunkReason); + } return success(); } @@ -1455,42 +1599,52 @@ checkSupportedStoreShape(VMIVRegType type, Value destination, buildWriteAccessPlan(destination, destinationType, type, VMIMemoryWriteMaskKind::AllTrue); if (!accessPlan.layoutSupport.isSupported()) { - if (reason) + if (reason) { *reason = accessPlan.layoutSupport.reason; + } return failure(); } - if (failed(checkSupportedMaskableVReg(type, reason))) + if (failed(checkSupportedMaskableVReg(type, reason))) { return failure(); + } VMILayoutSupport supports; - if (failed(supports.getStoreLayoutFact(type, reason))) + if (failed(supports.getStoreLayoutFact(type, reason))) { return failure(); + } - if (getDenseLaneStrideStoreDistToken(type)) + if (getDenseLaneStrideStoreDistToken(type)) { return success(); + } std::string fullChunkReason; - if (succeeded(checkFullDataPhysicalChunks(type, &fullChunkReason))) + if (succeeded(checkFullDataPhysicalChunks(type, &fullChunkReason))) { return success(); + } auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; VMILayoutAttr layout = type.getLayoutAttr(); - if (!layout) + if (!layout) { return fail("requires assigned layout"); - if (failed(getDataLanesPerPart(type.getElementType()))) + } + if (failed(getDataLanesPerPart(type.getElementType()))) { return fail("requires known physical lanes per part"); - if (layout.isContiguous() && layout.getLaneStride() == 1) + } + if (layout.isContiguous() && layout.getLaneStride() == 1) { return success(); + } std::string materializationReason; - if (succeeded(checkCanMaterializeToContiguous(type, &materializationReason))) + if (succeeded(checkCanMaterializeToContiguous(type, &materializationReason))) { return success(); + } return fail(Twine("partial/tail store requires contiguous layout or " "deinterleaved layout that can materialize to contiguous; " "value ") + @@ -1501,8 +1655,9 @@ LogicalResult checkSupportedInterleaveStoreShape( VMIInterleaveStoreOp op, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -1516,19 +1671,23 @@ LogicalResult checkSupportedInterleaveStoreShape( if (lowType.getElementCount() != highType.getElementCount() || lowType.getElementType() != highType.getElementType()) return fail("requires matching low/high input shape and element type"); - if (!getX2MemoryDistToken(lowType.getElementType(), "INTLV")) + if (!getX2MemoryDistToken(lowType.getElementType(), "INTLV")) { return fail("requires 8/16/32-bit element type for vstsx2 INTLV"); + } VMIMemoryAccessPlan accessPlan = buildWriteAccessPlan(op.getDestination(), op.getDestination().getType(), lowType, VMIMemoryWriteMaskKind::AllTrue); - if (!accessPlan.layoutSupport.isSupported()) + if (!accessPlan.layoutSupport.isSupported()) { return fail(accessPlan.layoutSupport.reason); - if (failed(checkSupportedMaskableVReg(lowType, reason))) + } + if (failed(checkSupportedMaskableVReg(lowType, reason))) { return failure(); + } std::string fullChunkReason; - if (failed(checkFullDataPhysicalChunks(lowType, &fullChunkReason))) + if (failed(checkFullDataPhysicalChunks(lowType, &fullChunkReason))) { return fail(Twine("requires full physical chunks; ") + fullChunkReason); + } return success(); } @@ -1536,34 +1695,41 @@ FailureOr getGroupSizeFromNumGroups(VMIVRegType type, int64_t numGroups, std::string *reason = nullptr) { auto fail = [&](const Twine &message) -> FailureOr { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; - if (numGroups <= 0) + if (numGroups <= 0) { return fail("requires num_groups to be positive"); - if (type.getElementCount() % numGroups != 0) + } + if (type.getElementCount() % numGroups != 0) { return fail("requires num_groups to evenly divide logical lane count"); + } return type.getElementCount() / numGroups; } LogicalResult checkSupportedGroupChunkShape(VMIVRegType type, int64_t groupSize, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; VMILayoutAttr layout = type.getLayoutAttr(); - if (!layout || !layout.isContiguous()) + if (!layout || !layout.isContiguous()) { return fail("requires assigned contiguous layout"); + } std::string fullChunkReason; - if (failed(checkFullDataPhysicalChunks(type, &fullChunkReason))) + if (failed(checkFullDataPhysicalChunks(type, &fullChunkReason))) { return fail(Twine("requires full physical chunks; ") + fullChunkReason); + } FailureOr lanesPerPart = getDataLanesPerPart(type.getElementType()); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return fail("requires known physical lanes per part"); + } if (groupSize <= 0 || type.getElementCount() % groupSize != 0) return fail("requires derived group size to evenly divide logical lane " "count"); @@ -1578,8 +1744,9 @@ LogicalResult checkDeinterleaved2GroupStoreChunkShape( int64_t *groupCount, int64_t *chunksPerGroupPerPart, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -1588,19 +1755,23 @@ LogicalResult checkDeinterleaved2GroupStoreChunkShape( layout.getLaneStride() != 1) return fail("requires deinterleaved=2 value layout"); std::string fullChunkReason; - if (failed(checkFullDataPhysicalChunks(type, &fullChunkReason))) + if (failed(checkFullDataPhysicalChunks(type, &fullChunkReason))) { return fail(Twine("requires full physical chunks; ") + fullChunkReason); + } FailureOr lanes = getDataLanesPerPart(type.getElementType()); - if (failed(lanes)) + if (failed(lanes)) { return fail("requires known physical lanes per part"); - if (!getX2MemoryDistToken(type.getElementType(), "INTLV")) + } + if (!getX2MemoryDistToken(type.getElementType(), "INTLV")) { return fail("requires 8/16/32-bit element type for vstsx2 INTLV"); + } if (groupSize <= 0 || type.getElementCount() % groupSize != 0) return fail("requires derived group size to evenly divide logical lane " "count"); int64_t pairLanes = 2 * *lanes; - if (groupSize % pairLanes != 0) + if (groupSize % pairLanes != 0) { return fail("requires group size to be a multiple of two physical chunks"); + } FailureOr part0Chunks = getDataChunksInPart(type, /*part=*/0); FailureOr part1Chunks = getDataChunksInPart(type, /*part=*/1); @@ -1611,51 +1782,59 @@ LogicalResult checkDeinterleaved2GroupStoreChunkShape( *lanesPerPart = *lanes; *groupCount = type.getElementCount() / groupSize; *chunksPerGroupPerPart = groupSize / pairLanes; - if (*part0Chunks != *groupCount * *chunksPerGroupPerPart) + if (*part0Chunks != *groupCount * *chunksPerGroupPerPart) { return fail("requires deinterleaved chunks to align with group rows"); + } return success(); } LogicalResult checkSupportedGroupLoadShape(VMIGroupLoadOp op, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; auto resultType = cast(op.getResult().getType()); VMILayoutAttr resultLayout = resultType.getLayoutAttr(); - if (!resultLayout) + if (!resultLayout) { return fail("requires assigned result layout"); + } FailureOr groupSize = getGroupSizeFromNumGroups( resultType, op.getNumGroupsAttr().getInt(), reason); - if (failed(groupSize)) + if (failed(groupSize)) { return failure(); + } if (resultLayout.isContiguous()) { VMILayoutSupport supports; - if (failed(supports.getGroupLoadLayoutFact(op, reason))) + if (failed(supports.getGroupLoadLayoutFact(op, reason))) { return failure(); + } if (failed(checkSupportedLoadShape(resultType, op.getSource(), op.getSource().getType(), std::nullopt, reason))) return failure(); std::optional rowStride = getConstantIndexValue(op.getRowStride()); - if (rowStride && *rowStride == *groupSize) + if (rowStride && *rowStride == *groupSize) { return success(); + } return checkSupportedGroupChunkShape(resultType, *groupSize, reason); } if (resultLayout.isBlockDeinterleaved() && resultType.getElementType().isF32()) { VMILayoutSupport supports; - if (failed(supports.getGroupLoadLayoutFact(op, reason))) + if (failed(supports.getGroupLoadLayoutFact(op, reason))) { return failure(); + } VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(op.getSource(), op.getSource().getType(), resultType, getConstantIndexValue(op.getOffset()), VMIMemoryValidMaskKind::AllTrue); - if (!accessPlan.layoutSupport.isSupported()) + if (!accessPlan.layoutSupport.isSupported()) { return fail(accessPlan.layoutSupport.reason); + } if (!isa(op.getSource().getType())) return fail( "block_deinterleaved group_load requires !pto.ptr source"); @@ -1682,8 +1861,9 @@ LogicalResult checkSupportedGroupSlotLoadShape( VMIGroupSlotLoadOp op, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -1691,15 +1871,18 @@ LogicalResult checkSupportedGroupSlotLoadShape( VMILayoutSupport supports; FailureOr fact = supports.getGroupSlotLoadLayoutFact( resultType, op.getNumGroupsAttr().getInt(), reason); - if (failed(fact)) + if (failed(fact)) { return failure(); + } VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(op.getSource(), op.getSource().getType(), resultType, getConstantIndexValue(op.getOffset()), VMIMemoryValidMaskKind::AllTrue); - if (!accessPlan.layoutSupport.isSupported()) + if (!accessPlan.layoutSupport.isSupported()) { return fail(accessPlan.layoutSupport.reason); - if (!isa(op.getSource().getType())) + } + if (!isa(op.getSource().getType())) { return fail("group_slot_load requires !pto.ptr source"); + } if (fact->slots == 8) { std::optional sourceGroupStride = @@ -1712,8 +1895,9 @@ LogicalResult checkSupportedGroupSlotLoadShape( unsigned elementBits = pto::getPTOStorageElemBitWidth(resultType.getElementType()); - if (elementBits == 0 || 256 % elementBits != 0) + if (elementBits == 0 || 256 % elementBits != 0) { return fail("slots=1 group_slot_load requires supported element width"); + } int64_t alignedStrideElems = 256 / elementBits; std::optional sourceGroupStride = getConstantIndexValue(op.getSourceGroupStride()); @@ -1732,21 +1916,25 @@ LogicalResult checkSupportedGroupBroadcastLoadShape( VMIGroupBroadcastLoadOp op, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; VMILayoutSupport supports; - if (failed(supports.getGroupBroadcastLoadSupport(op, reason))) + if (failed(supports.getGroupBroadcastLoadSupport(op, reason))) { return failure(); + } VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(op.getSource(), op.getSource().getType(), cast(op.getResult().getType()), getConstantIndexValue(op.getOffset()), VMIMemoryValidMaskKind::AllTrue); - if (!accessPlan.layoutSupport.isSupported()) + if (!accessPlan.layoutSupport.isSupported()) { return fail(accessPlan.layoutSupport.reason); - if (!isa(op.getSource().getType())) + } + if (!isa(op.getSource().getType())) { return fail("group_broadcast_load requires !pto.ptr source"); + } return success(); } @@ -1779,8 +1967,9 @@ FailureOr getOneBlockGroupStorePlan( const VMIGroupStoreLayoutFact &fact, std::string *reason) { auto fail = [&](const Twine &message) -> FailureOr { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -1792,8 +1981,9 @@ FailureOr getOneBlockGroupStorePlan( fact.lanesPerPart <= 0 || fact.lanesPerPart % fact.groupSize != 0) return fail("one-block group_store requires one 32B group per VCG block"); - if (!isa(op.getDestination().getType())) + if (!isa(op.getDestination().getType())) { return fail("one-block group_store requires !pto.ptr destination"); + } std::optional rowStride = getConstantIndexValue(op.getRowStride()); @@ -1813,8 +2003,9 @@ FailureOr getOneBlockGroupStorePlan( LogicalResult checkSupportedGroupStoreShape(VMIGroupStoreOp op, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -1823,32 +2014,37 @@ checkSupportedGroupStoreShape(VMIGroupStoreOp op, std::string *reason) { std::optional rowStride = getConstantIndexValue(op.getRowStride()); if (isCompactSmallGroupStore(layout, valueType, op.getNumGroupsAttr().getInt(), rowStride)) { - if (!isa(op.getDestination().getType())) + if (!isa(op.getDestination().getType())) { return fail("compact small group_store requires !pto.ptr destination"); + } VMIMemoryAccessPlan accessPlan = buildWriteAccessPlan( op.getDestination(), op.getDestination().getType(), valueType, VMIMemoryWriteMaskKind::AllTrue); - if (!accessPlan.layoutSupport.isSupported()) + if (!accessPlan.layoutSupport.isSupported()) { return fail(accessPlan.layoutSupport.reason); + } return success(); } if (layout && layout.isGroupSlots()) { VMILayoutSupport supports; FailureOr fact = supports.getGroupStoreLayoutFact( valueType, op.getNumGroupsAttr().getInt(), reason); - if (failed(fact)) + if (failed(fact)) { return failure(); + } VMIMemoryAccessPlan accessPlan = buildWriteAccessPlan(op.getDestination(), op.getDestination().getType(), valueType, VMIMemoryWriteMaskKind::AllTrue); - if (!accessPlan.layoutSupport.isSupported()) + if (!accessPlan.layoutSupport.isSupported()) { return fail(accessPlan.layoutSupport.reason); + } if (fact->slots == 1) { unsigned elementBits = pto::getPTOStorageElemBitWidth(valueType.getElementType()); - if (elementBits == 0 || 256 % elementBits != 0) + if (elementBits == 0 || 256 % elementBits != 0) { return fail("slots=1 group_store requires supported element width"); + } std::optional rowStride = getConstantIndexValue(op.getRowStride()); if (rowStride && *rowStride <= 0) @@ -1869,15 +2065,17 @@ checkSupportedGroupStoreShape(VMIGroupStoreOp op, std::string *reason) { VMILayoutSupport supports; FailureOr fact = supports.getGroupStoreLayoutFact(op, valueType, reason); - if (failed(fact)) + if (failed(fact)) { return failure(); + } if (failed(checkSupportedStoreShape(valueType, op.getDestination(), op.getDestination().getType(), reason))) return failure(); if (fact->blockClass == VMIGroupBlockClass::OneBlock) { - if (failed(getOneBlockGroupStorePlan(op, valueType, *fact, reason))) + if (failed(getOneBlockGroupStorePlan(op, valueType, *fact, reason))) { return failure(); + } return success(); } if (succeeded( @@ -1895,8 +2093,9 @@ checkSupportedGroupStoreShape(VMIGroupStoreOp op, std::string *reason) { LogicalResult checkSupportedMaskedLoadShape(VMIMaskedLoadOp op, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -1909,20 +2108,24 @@ checkSupportedMaskedLoadShape(VMIMaskedLoadOp op, std::string *reason) { VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(op.getSource(), op.getSource().getType(), resultType, getConstantIndexValue(op.getOffset()), VMIMemoryValidMaskKind::ExplicitMask); - if (!accessPlan.layoutSupport.isSupported()) + if (!accessPlan.layoutSupport.isSupported()) { return fail(accessPlan.layoutSupport.reason); - if (!resultLayout || !passthruLayout || !maskLayout) + } + if (!resultLayout || !passthruLayout || !maskLayout) { return fail("requires assigned result, passthru, and mask layouts"); + } if (!resultLayout.isContiguous() || !passthruLayout.isContiguous() || !maskLayout.isContiguous()) return fail("requires contiguous result, passthru, and mask layouts"); std::string fullChunkReason; - if (succeeded(checkFullDataPhysicalChunks(resultType, &fullChunkReason))) + if (succeeded(checkFullDataPhysicalChunks(resultType, &fullChunkReason))) { return success(); + } - if (accessPlan.safeReadProof.proven) + if (accessPlan.safeReadProof.proven) { return success(); + } requireUnavailableReadFallback(accessPlan); return fail(Twine("partial/tail masked_load requires statically safe " "full-read footprint; value ") + @@ -1934,8 +2137,9 @@ checkSupportedMaskedLoadShape(VMIMaskedLoadOp op, std::string *reason) { LogicalResult checkSupportedGatherShape(VMIGatherOp op, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -1962,8 +2166,9 @@ checkSupportedGatherShape(VMIGatherOp op, std::string *reason) { unsigned resultBits = pto::getPTOStorageElemBitWidth(resultType.getElementType()); auto indexElementType = dyn_cast(indicesType.getElementType()); - if (!indexElementType || indexElementType.isSigned()) + if (!indexElementType || indexElementType.isSigned()) { return fail("requires signless or unsigned integer indices"); + } bool isU16Gather = resultBits == 16 && indexElementType.isUnsigned() && indexElementType.getWidth() == 16 && maskType.getGranularity() == "b16"; @@ -1999,8 +2204,9 @@ checkSupportedGatherShape(VMIGatherOp op, std::string *reason) { if (failed(checkFullDataPhysicalChunks(passthruType, &passthruReason))) return fail(Twine("passthru requires full physical chunks; ") + passthruReason); - if (failed(checkFullVMIPhysicalChunks(maskType, &maskReason))) + if (failed(checkFullVMIPhysicalChunks(maskType, &maskReason))) { return fail(Twine("mask requires full physical chunks; ") + maskReason); + } } else if (*resultArity != 1) { return fail("ui16 gather currently supports one physical chunk"); } @@ -2011,8 +2217,9 @@ checkSupportedGatherShape(VMIGatherOp op, std::string *reason) { LogicalResult checkSupportedScatterShape(VMIScatterOp op, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -2022,8 +2229,9 @@ checkSupportedScatterShape(VMIScatterOp op, std::string *reason) { VMILayoutAttr valueLayout = valueType.getLayoutAttr(); VMILayoutAttr indicesLayout = indicesType.getLayoutAttr(); VMILayoutAttr maskLayout = maskType.getLayoutAttr(); - if (!valueLayout || !indicesLayout || !maskLayout) + if (!valueLayout || !indicesLayout || !maskLayout) { return fail("requires assigned value, indices, and mask layouts"); + } if (!valueLayout.isContiguous() || !indicesLayout.isContiguous() || !maskLayout.isContiguous()) return fail("requires contiguous value, indices, and mask layouts"); @@ -2035,8 +2243,9 @@ checkSupportedScatterShape(VMIScatterOp op, std::string *reason) { unsigned valueBits = pto::getPTOStorageElemBitWidth(valueType.getElementType()); auto indexElementType = dyn_cast(indicesType.getElementType()); - if (!indexElementType || indexElementType.isSigned()) + if (!indexElementType || indexElementType.isSigned()) { return fail("requires signless or unsigned integer indices"); + } bool isB8Scatter = valueBits == 8 && indexElementType.getWidth() == 16 && maskType.getGranularity() == "b16"; @@ -2054,8 +2263,9 @@ checkSupportedScatterShape(VMIScatterOp op, std::string *reason) { FailureOr valueArity = getVMIPhysicalArity(valueType); FailureOr indicesArity = getVMIPhysicalArity(indicesType); FailureOr maskArity = getVMIPhysicalArity(maskType); - if (failed(valueArity) || failed(indicesArity) || failed(maskArity)) + if (failed(valueArity) || failed(indicesArity) || failed(maskArity)) { return fail("requires computable physical arity"); + } if (*valueArity != *indicesArity || *valueArity != *maskArity) return fail("requires value, indices, and mask to have the same physical " "arity"); @@ -2063,13 +2273,15 @@ checkSupportedScatterShape(VMIScatterOp op, std::string *reason) { std::string valueReason; std::string indicesReason; std::string maskReason; - if (failed(checkFullDataPhysicalChunks(valueType, &valueReason))) + if (failed(checkFullDataPhysicalChunks(valueType, &valueReason))) { return fail(Twine("value requires full physical chunks; ") + valueReason); + } if (failed(checkFullDataPhysicalChunks(indicesType, &indicesReason))) return fail(Twine("indices require full physical chunks; ") + indicesReason); - if (failed(checkFullVMIPhysicalChunks(maskType, &maskReason))) + if (failed(checkFullVMIPhysicalChunks(maskType, &maskReason))) { return fail(Twine("mask requires full physical chunks; ") + maskReason); + } return success(); } @@ -2077,8 +2289,9 @@ checkSupportedScatterShape(VMIScatterOp op, std::string *reason) { LogicalResult checkSupportedStrideStoreShape(VMIStrideStoreOp op, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -2086,10 +2299,12 @@ checkSupportedStrideStoreShape(VMIStrideStoreOp op, std::string *reason) { auto maskType = cast(op.getMask().getType()); VMILayoutAttr valueLayout = valueType.getLayoutAttr(); VMILayoutAttr maskLayout = maskType.getLayoutAttr(); - if (!valueLayout || !maskLayout) + if (!valueLayout || !maskLayout) { return fail("requires assigned value and mask layouts"); - if (!valueLayout.isContiguous() || !maskLayout.isContiguous()) + } + if (!valueLayout.isContiguous() || !maskLayout.isContiguous()) { return fail("requires contiguous value and mask layouts"); + } if (!isa(op.getDestination().getType())) return fail("requires !pto.ptr destination because pto.vsstb is " @@ -2101,18 +2316,21 @@ checkSupportedStrideStoreShape(VMIStrideStoreOp op, std::string *reason) { FailureOr valueArity = getVMIPhysicalArity(valueType); FailureOr maskArity = getVMIPhysicalArity(maskType); - if (failed(valueArity) || failed(maskArity)) + if (failed(valueArity) || failed(maskArity)) { return fail("requires computable physical arity"); - if (*valueArity != 1 || *maskArity != 1) + } + if (*valueArity != 1 || *maskArity != 1) { return fail("currently supports one physical value/mask chunk"); + } return success(); } LogicalResult checkSupportedStrideLoadShape(VMIStrideLoadOp op, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -2120,20 +2338,25 @@ checkSupportedStrideLoadShape(VMIStrideLoadOp op, std::string *reason) { auto maskType = cast(op.getMask().getType()); VMILayoutAttr resultLayout = resultType.getLayoutAttr(); VMILayoutAttr maskLayout = maskType.getLayoutAttr(); - if (!resultLayout || !maskLayout) + if (!resultLayout || !maskLayout) { return fail("requires assigned result and mask layouts"); - if (!resultLayout.isContiguous() || !maskLayout.isContiguous()) + } + if (!resultLayout.isContiguous() || !maskLayout.isContiguous()) { return fail("requires contiguous result and mask layouts"); + } - if (!isa(op.getSource().getType())) + if (!isa(op.getSource().getType())) { return fail("requires !pto.ptr source because pto.vsldb is pointer-only"); + } FailureOr resultArity = getVMIPhysicalArity(resultType); FailureOr maskArity = getVMIPhysicalArity(maskType); - if (failed(resultArity) || failed(maskArity)) + if (failed(resultArity) || failed(maskArity)) { return fail("requires computable physical arity"); - if (*resultArity != 1 || *maskArity != 1) + } + if (*resultArity != 1 || *maskArity != 1) { return fail("currently supports one physical result/mask chunk"); + } return success(); } @@ -2155,19 +2378,22 @@ bool isStaticAllActiveMask(Value mask, int64_t expectedLanes, std::string *reason = nullptr) { mask = stripMaskMaterialization(mask); auto fail = [&](const Twine &message) { - if (reason) + if (reason) { *reason = message.str(); + } return false; }; if (auto createMask = mask.getDefiningOp()) { auto activeConstant = createMask.getActiveLanes().getDefiningOp(); - if (!activeConstant) + if (!activeConstant) { return fail("create_mask active_lanes is dynamic"); + } auto activeAttr = dyn_cast(activeConstant.getValue()); - if (!activeAttr) + if (!activeAttr) { return fail("create_mask active_lanes is not an integer constant"); + } return activeAttr.getInt() >= expectedLanes ? true : fail("create_mask active_lanes is smaller than the logical " @@ -2176,15 +2402,17 @@ bool isStaticAllActiveMask(Value mask, int64_t expectedLanes, if (auto constantMask = mask.getDefiningOp()) { auto denseAttr = dyn_cast(constantMask.getValue()); - if (!denseAttr) + if (!denseAttr) { return fail("constant_mask is not a dense integer mask"); + } if (denseAttr.getNumElements() != expectedLanes) return fail("constant_mask element count does not match the logical " "lane count"); auto values = denseAttr.getValues(); for (bool value : values) - if (!value) + if (!value) { return fail("constant_mask contains an inactive lane"); + } return true; } @@ -2194,8 +2422,9 @@ bool isStaticAllActiveMask(Value mask, int64_t expectedLanes, LogicalResult checkSupportedExpandLoadShape(VMIExpandLoadOp op, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -2208,10 +2437,12 @@ checkSupportedExpandLoadShape(VMIExpandLoadOp op, std::string *reason) { VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(op.getSource(), op.getSource().getType(), resultType, getConstantIndexValue(op.getOffset()), VMIMemoryValidMaskKind::ExplicitMask); - if (!accessPlan.layoutSupport.isSupported()) + if (!accessPlan.layoutSupport.isSupported()) { return fail(accessPlan.layoutSupport.reason); - if (!resultLayout || !passthruLayout || !maskLayout) + } + if (!resultLayout || !passthruLayout || !maskLayout) { return fail("requires assigned result, passthru, and mask layouts"); + } if (!resultLayout.isContiguous() || !passthruLayout.isContiguous() || !maskLayout.isContiguous()) return fail("requires contiguous result, passthru, and mask layouts"); @@ -2225,8 +2456,9 @@ checkSupportedExpandLoadShape(VMIExpandLoadOp op, std::string *reason) { succeeded(checkFullDataPhysicalChunks(resultType, &fullChunkReason))) return success(); - if (staticAllActive && accessPlan.safeReadProof.proven) + if (staticAllActive && accessPlan.safeReadProof.proven) { return success(); + } std::string allActivePathReason; if (!staticAllActive) { @@ -2250,14 +2482,16 @@ checkSupportedExpandLoadShape(VMIExpandLoadOp op, std::string *reason) { if (pto::getPTOStorageElemBitWidth(resultType.getElementType()) != 32) return fail("runtime-mask path currently requires 32-bit result element " "type so prefix indices and gather result lane counts match"); - if (maskType.getGranularity() != "b32") + if (maskType.getGranularity() != "b32") { return fail("runtime-mask path requires b32 mask granularity"); + } FailureOr resultArity = getVMIPhysicalArity(resultType); FailureOr passthruArity = getVMIPhysicalArity(passthruType); FailureOr maskArity = getVMIPhysicalArity(maskType); - if (failed(resultArity) || failed(passthruArity) || failed(maskArity)) + if (failed(resultArity) || failed(passthruArity) || failed(maskArity)) { return fail("runtime-mask path requires computable physical arity"); + } if (*resultArity != 1 || *passthruArity != 1 || *maskArity != 1) return fail("runtime-mask path currently supports only one physical " "chunk because prefix indices must not reset across chunks"); @@ -2285,8 +2519,9 @@ checkSupportedMaskedStoreShape(VMIVRegType valueType, VMIMaskType maskType, buildWriteAccessPlan(destination, destinationType, valueType, VMIMemoryWriteMaskKind::ExplicitMask); if (!accessPlan.layoutSupport.isSupported()) { - if (reason) + if (reason) { *reason = accessPlan.layoutSupport.reason; + } return failure(); } @@ -2297,20 +2532,23 @@ checkSupportedMaskedStoreShape(VMIVRegType valueType, VMIMaskType maskType, return success(); auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; VMILayoutAttr valueLayout = valueType.getLayoutAttr(); VMILayoutAttr maskLayout = maskType.getLayoutAttr(); - if (!valueLayout || !maskLayout) + if (!valueLayout || !maskLayout) { return fail("requires assigned value and mask layouts"); + } FailureOr valueArity = getVMIPhysicalArity(valueType); FailureOr maskArity = getVMIPhysicalArity(maskType); - if (failed(valueArity) || failed(maskArity) || *valueArity != *maskArity) + if (failed(valueArity) || failed(maskArity) || *valueArity != *maskArity) { return fail("requires matching value/mask physical arity"); + } if (valueLayout.hasDenseLaneStride()) { VMILayoutSupport supports; @@ -2343,8 +2581,9 @@ FailureOr getContiguousActiveDataLanes(VMIVRegType vmiType, int64_t chunk) { FailureOr lanesPerPart = getDataLanesPerPart(vmiType.getElementType()); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return failure(); + } int64_t remaining = vmiType.getElementCount() - chunk * *lanesPerPart; return std::clamp(remaining, 0, *lanesPerPart); @@ -2354,16 +2593,19 @@ FailureOr getActiveDataLanesInPhysicalChunk(VMIVRegType vmiType, int64_t chunk) { FailureOr lanesPerPart = getDataLanesPerPart(vmiType.getElementType()); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return failure(); + } int64_t active = 0; for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { FailureOr padding = isPaddingLane(vmiType, /*part=*/0, chunk, lane); - if (failed(padding)) + if (failed(padding)) { return failure(); - if (!*padding) + } + if (!*padding) { ++active; + } } return active; } @@ -2373,23 +2615,28 @@ FailureOr createContiguousStoreMask(Location loc, VMIVRegType vmiType, PatternRewriter &rewriter) { FailureOr lanesPerPart = getDataLanesPerPart(vmiType.getElementType()); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return failure(); + } FailureOr activeLanes = getContiguousActiveDataLanes(vmiType, chunk); - if (failed(activeLanes)) + if (failed(activeLanes)) { return failure(); - if (*activeLanes == *lanesPerPart) + } + if (*activeLanes == *lanesPerPart) { return createAllTrueMaskForVReg(loc, vregType, rewriter); + } FailureOr maskType = getMaskTypeForVReg(vregType, rewriter.getContext()); - if (failed(maskType)) + if (failed(maskType)) { return failure(); + } FailureOr> maskAndRemaining = createRuntimePrefixMask( loc, *maskType, createI32Constant(loc, *activeLanes, rewriter), rewriter); - if (failed(maskAndRemaining)) + if (failed(maskAndRemaining)) { return failure(); + } return maskAndRemaining->first; } @@ -2399,23 +2646,28 @@ FailureOr createMaskedStorePredicate(Location loc, VMIVRegType vmiType, PatternRewriter &rewriter) { FailureOr lanesPerPart = getDataLanesPerPart(vmiType.getElementType()); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return failure(); + } FailureOr activeLanes = getContiguousActiveDataLanes(vmiType, chunk); - if (failed(activeLanes)) + if (failed(activeLanes)) { return failure(); - if (*activeLanes == *lanesPerPart) + } + if (*activeLanes == *lanesPerPart) { return userMask; + } auto maskType = dyn_cast(userMask.getType()); - if (!maskType) + if (!maskType) { return failure(); + } FailureOr tailMask = createContiguousStoreMask(loc, vmiType, chunk, vregType, rewriter); FailureOr allTrue = createAllTrueMask(loc, maskType, rewriter); - if (failed(tailMask) || failed(allTrue)) + if (failed(tailMask) || failed(allTrue)) { return failure(); + } return rewriter.create(loc, maskType, userMask, *tailMask, *allTrue) .getResult(); } @@ -2424,13 +2676,15 @@ FailureOr createDenseLaneStrideStorePredicate( Location loc, VMIVRegType vmiType, int64_t chunk, Value userMask, StringRef targetGranularity, PatternRewriter &rewriter) { auto sourceMaskType = dyn_cast(userMask.getType()); - if (!sourceMaskType) + if (!sourceMaskType) { return failure(); + } auto targetMaskType = MaskType::get(rewriter.getContext(), targetGranularity); Value compactMask = userMask; VMILayoutAttr layout = vmiType.getLayoutAttr(); - if (!layout) + if (!layout) { return failure(); + } auto lower = rewriter.getStringAttr("LOWER"); StringRef sourceGranularity = sourceMaskType.getGranularity(); @@ -2456,16 +2710,19 @@ FailureOr createDenseLaneStrideStorePredicate( FailureOr activeLanes = getActiveDataLanesInPhysicalChunk(vmiType, chunk); FailureOr maskLanes = getMaskLanesPerPart(targetGranularity); - if (failed(activeLanes) || failed(maskLanes)) + if (failed(activeLanes) || failed(maskLanes)) { return failure(); - if (*activeLanes == *maskLanes) + } + if (*activeLanes == *maskLanes) { return compactMask; + } FailureOr tailMask = createPrefixMaskForActiveLanes( loc, targetMaskType, *activeLanes, rewriter); FailureOr allTrue = createAllTrueMask(loc, targetMaskType, rewriter); - if (failed(tailMask) || failed(allTrue)) + if (failed(tailMask) || failed(allTrue)) { return failure(); + } return rewriter .create(loc, targetMaskType, compactMask, *tailMask, *allTrue) .getResult(); @@ -2474,8 +2731,9 @@ FailureOr createDenseLaneStrideStorePredicate( FailureOr> computeShuffleForwardingSourceParts(VMIShuffleOp op, std::string *reason) { auto fail = [&](const Twine &message) -> FailureOr> { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -2483,23 +2741,27 @@ computeShuffleForwardingSourceParts(VMIShuffleOp op, std::string *reason) { auto resultType = cast(op.getResult().getType()); FailureOr lanesPerPart = getDataLanesPerPart(sourceType.getElementType()); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return fail("requires known lanes per physical part"); + } ArrayRef indices = op.getIndices(); - if (indices.empty()) + if (indices.empty()) { return fail("requires non-empty indices"); + } FailureOr resultFactor = getDataLayoutFactor(resultType); - if (failed(resultFactor)) + if (failed(resultFactor)) { return fail("requires assigned result layout"); + } SmallVector sourceFlatIndices; for (int64_t resultPart = 0; resultPart < *resultFactor; ++resultPart) { FailureOr resultChunks = getDataChunksInPart(resultType, resultPart); - if (failed(resultChunks)) + if (failed(resultChunks)) { return fail("requires known result physical chunks"); + } for (int64_t resultChunk = 0; resultChunk < *resultChunks; ++resultChunk) { std::optional sourcePart; @@ -2507,10 +2769,12 @@ computeShuffleForwardingSourceParts(VMIShuffleOp op, std::string *reason) { for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { FailureOr padding = isPaddingLane(resultType, resultPart, resultChunk, lane); - if (failed(padding)) + if (failed(padding)) { return fail("failed to classify result padding lanes"); - if (*padding) + } + if (*padding) { continue; + } FailureOr resultLogicalLane = mapPhysicalLaneToLogical(resultType, resultPart, resultChunk, lane); @@ -2520,10 +2784,12 @@ computeShuffleForwardingSourceParts(VMIShuffleOp op, std::string *reason) { FailureOr sourcePhysical = mapLogicalLaneToPhysical(sourceType, indices[*resultLogicalLane]); - if (failed(sourcePhysical)) + if (failed(sourcePhysical)) { return fail("failed to map source lane"); - if (sourcePhysical->lane != lane) + } + if (sourcePhysical->lane != lane) { return fail("requires same-lane physical chunks"); + } if (!sourcePart) { sourcePart = sourcePhysical->part; @@ -2535,12 +2801,14 @@ computeShuffleForwardingSourceParts(VMIShuffleOp op, std::string *reason) { return fail("requires one source chunk per result chunk"); } - if (!sourcePart || !sourceChunk) + if (!sourcePart || !sourceChunk) { return fail("requires at least one logical lane per result chunk"); + } FailureOr sourceFlatIndex = getDataFlatPartIndex(sourceType, *sourcePart, *sourceChunk); - if (failed(sourceFlatIndex)) + if (failed(sourceFlatIndex)) { return fail("source part range is out of bounds"); + } sourceFlatIndices.push_back(*sourceFlatIndex); } } @@ -2557,26 +2825,30 @@ struct ShuffleVselrPlan { FailureOr computeShuffleLane0SplatSourcePart(VMIShuffleOp op, std::string *reason) { auto fail = [&](const Twine &message) -> FailureOr { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; ArrayRef indices = op.getIndices(); - if (indices.empty()) + if (indices.empty()) { return fail("requires non-empty indices"); + } if (!llvm::all_of(indices, [](int64_t index) { return index == 0; })) return fail("requires every result lane to select source lane 0"); auto sourceType = cast(op.getSource().getType()); FailureOr sourceLane = mapLogicalLaneToPhysical(sourceType, 0); - if (failed(sourceLane)) + if (failed(sourceLane)) { return fail("failed to map source lane 0"); + } FailureOr sourceFlatIndex = getDataFlatPartIndex(sourceType, sourceLane->part, sourceLane->chunk); - if (failed(sourceFlatIndex)) + if (failed(sourceFlatIndex)) { return fail("source lane 0 part range is out of bounds"); + } return *sourceFlatIndex; } @@ -2584,8 +2856,9 @@ FailureOr> computeShuffleVselrPlans(VMIShuffleOp op, std::string *reason) { auto fail = [&](const Twine &message) -> FailureOr> { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -2593,23 +2866,27 @@ computeShuffleVselrPlans(VMIShuffleOp op, std::string *reason) { auto resultType = cast(op.getResult().getType()); FailureOr lanesPerPart = getDataLanesPerPart(sourceType.getElementType()); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return fail("requires known lanes per physical part"); + } ArrayRef indices = op.getIndices(); - if (indices.empty()) + if (indices.empty()) { return fail("requires non-empty indices"); + } FailureOr resultFactor = getDataLayoutFactor(resultType); - if (failed(resultFactor)) + if (failed(resultFactor)) { return fail("requires assigned result layout"); + } SmallVector plans; for (int64_t resultPart = 0; resultPart < *resultFactor; ++resultPart) { FailureOr resultChunks = getDataChunksInPart(resultType, resultPart); - if (failed(resultChunks)) + if (failed(resultChunks)) { return fail("requires known result physical chunks"); + } for (int64_t resultChunk = 0; resultChunk < *resultChunks; ++resultChunk) { std::optional sourcePart; @@ -2619,8 +2896,9 @@ computeShuffleVselrPlans(VMIShuffleOp op, std::string *reason) { for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { FailureOr padding = isPaddingLane(resultType, resultPart, resultChunk, lane); - if (failed(padding) || *padding) + if (failed(padding) || *padding) { return fail("requires full physical result chunks"); + } FailureOr resultLogicalLane = mapPhysicalLaneToLogical(resultType, resultPart, resultChunk, lane); @@ -2630,8 +2908,9 @@ computeShuffleVselrPlans(VMIShuffleOp op, std::string *reason) { FailureOr sourcePhysical = mapLogicalLaneToPhysical(sourceType, indices[*resultLogicalLane]); - if (failed(sourcePhysical)) + if (failed(sourcePhysical)) { return fail("failed to map source lane"); + } if (!sourcePart) { sourcePart = sourcePhysical->part; @@ -2648,22 +2927,25 @@ computeShuffleVselrPlans(VMIShuffleOp op, std::string *reason) { int64_t descExpected = *baseLane - lane; bool asc = sourcePhysical->lane == ascExpected; bool desc = sourcePhysical->lane == descExpected; - if (!asc && !desc) + if (!asc && !desc) { return fail("requires ASC or DESC affine source lane indices"); + } bool laneDescending = desc && !asc; if (!descending) { descending = laneDescending; continue; } - if (*descending != laneDescending) + if (*descending != laneDescending) { return fail("requires one index order per result chunk"); + } } FailureOr sourceFlatIndex = getDataFlatPartIndex(sourceType, *sourcePart, *sourceChunk); - if (failed(sourceFlatIndex)) + if (failed(sourceFlatIndex)) { return fail("source part range is out of bounds"); + } plans.push_back(ShuffleVselrPlan{*sourceFlatIndex, *baseLane, descending.value_or(false)}); } @@ -2680,14 +2962,16 @@ FailureOr> computeConstantMaskMaterialization(VMIConstantMaskOp op, std::string *reason) { auto fail = [&](const Twine &message) -> FailureOr> { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; auto denseAttr = dyn_cast(op.getValue()); - if (!denseAttr) + if (!denseAttr) { return fail("only dense integer mask constants are supported"); + } auto resultVMIType = cast(op.getResult().getType()); VMILayoutAttr layout = resultVMIType.getLayoutAttr(); @@ -2701,8 +2985,9 @@ computeConstantMaskMaterialization(VMIConstantMaskOp op, std::string *reason) { failed(physicalGranularity) ? FailureOr(failure()) : getMaskLanesPerPart(*physicalGranularity); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return fail("requires known physical mask lanes per part"); + } auto boolValues = denseAttr.getValues(); int64_t factor = layout.isDenseSplit() ? layout.getFactor() : 1; @@ -2715,8 +3000,9 @@ computeConstantMaskMaterialization(VMIConstantMaskOp op, std::string *reason) { for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { FailureOr padding = isPaddingLane(resultVMIType, part, chunk, lane); - if (failed(padding)) + if (failed(padding)) { return fail("failed to map physical padding lane"); + } if (*padding) { materialization.activeLanes.push_back(0); continue; @@ -2725,12 +3011,14 @@ computeConstantMaskMaterialization(VMIConstantMaskOp op, std::string *reason) { FailureOr logicalLane = mapPhysicalLaneToLogical(resultVMIType, part, chunk, lane); - if (failed(logicalLane)) + if (failed(logicalLane)) { return fail("failed to map physical lane"); + } materialization.activeLanes.push_back(boolValues[*logicalLane] ? 1 : 0); } - if (!anyLane) + if (!anyLane) { break; + } materializations.push_back(std::move(materialization)); } } @@ -2744,18 +3032,21 @@ computeGroupMaskMaterializationForType(VMICreateGroupMaskOp op, std::string *reason) { auto fail = [&](const Twine &message) -> FailureOr> { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; auto activeConstant = op.getActiveElemsPerGroup().getDefiningOp(); - if (!activeConstant) + if (!activeConstant) { return fail("requires constant active_elems_per_group"); + } auto activeAttr = dyn_cast(activeConstant.getValue()); - if (!activeAttr) + if (!activeAttr) { return fail("active_elems_per_group must be an integer constant"); + } VMILayoutAttr layout = resultVMIType.getLayoutAttr(); if (!layout || @@ -2768,8 +3059,9 @@ computeGroupMaskMaterializationForType(VMICreateGroupMaskOp op, failed(physicalGranularity) ? FailureOr(failure()) : getMaskLanesPerPart(*physicalGranularity); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return fail("requires known physical mask lanes per part"); + } int64_t numGroups = op.getNumGroupsAttr().getInt(); int64_t groupSize = op.getGroupSizeAttr().getInt(); @@ -2778,10 +3070,12 @@ computeGroupMaskMaterializationForType(VMICreateGroupMaskOp op, return fail("requires result lane count to match num_groups * group_size"); int64_t activeElems = activeAttr.getInt(); - if (activeElems < 0) + if (activeElems < 0) { activeElems = 0; - if (activeElems > groupSize) + } + if (activeElems > groupSize) { activeElems = groupSize; + } int64_t factor = layout.isDenseSplit() ? layout.getFactor() : 1; SmallVector materializations; @@ -2793,8 +3087,9 @@ computeGroupMaskMaterializationForType(VMICreateGroupMaskOp op, for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { FailureOr padding = isPaddingLane(resultVMIType, part, chunk, lane); - if (failed(padding)) + if (failed(padding)) { return fail("failed to map physical padding lane"); + } if (*padding) { materialization.activeLanes.push_back(0); continue; @@ -2803,14 +3098,16 @@ computeGroupMaskMaterializationForType(VMICreateGroupMaskOp op, FailureOr logicalLane = mapPhysicalLaneToLogical(resultVMIType, part, chunk, lane); - if (failed(logicalLane)) + if (failed(logicalLane)) { return fail("failed to map physical lane"); + } int64_t laneInGroup = *logicalLane % groupSize; materialization.activeLanes.push_back(laneInGroup < activeElems ? 1 : 0); } - if (!anyLane) + if (!anyLane) { break; + } materializations.push_back(std::move(materialization)); } } @@ -2831,16 +3128,19 @@ FailureOr materializeConstantMaskChunk(Location loc, MaskType maskType, FailureOr createPowerOfTwoRemainder(Location loc, Value value, int64_t modulus, Value allMask, PatternRewriter &rewriter) { - if (modulus <= 0) + if (modulus <= 0) { return failure(); + } auto vectorType = dyn_cast(value.getType()); - if (!vectorType) + if (!vectorType) { return failure(); + } std::optional shift = getPowerOfTwoLog2(modulus); - if (!shift) + if (!shift) { return failure(); + } if (*shift == 0) { Value zero = createI32Constant(loc, 0, rewriter); return rewriter.create(loc, vectorType, zero, allMask, @@ -2869,10 +3169,12 @@ FailureOr> materializeDynamicGroupMaskForType( }; VMILayoutAttr layout = resultVMIType.getLayoutAttr(); - if (!layout) + if (!layout) { return fail("dynamic create_group_mask requires assigned layout"); - if (layout.getLaneStride() != 1) + } + if (layout.getLaneStride() != 1) { return fail("dynamic create_group_mask requires lane_stride=1 layout"); + } if (resultVMIType.getGranularity() != "b32") return fail("dynamic create_group_mask currently requires b32 " "granularity"); @@ -2894,8 +3196,9 @@ FailureOr> materializeDynamicGroupMaskForType( if (failed(lanesPerPart) || failed(arity) || *arity < 1) return fail("dynamic create_group_mask requires computable physical " "mask chunks"); - if (static_cast(resultTypes.size()) != *arity) + if (static_cast(resultTypes.size()) != *arity) { return fail("dynamic create_group_mask physical result count mismatch"); + } std::optional groupShift = getPowerOfTwoLog2(groupSize); if (!groupShift) @@ -2926,12 +3229,14 @@ FailureOr> materializeDynamicGroupMaskForType( for (int64_t chunk = 0; chunk < chunksPerPart; ++chunk) { Type resultType = resultTypes[part * chunksPerPart + chunk]; auto maskType = dyn_cast(resultType); - if (!maskType || !maskType.isB32()) + if (!maskType || !maskType.isB32()) { return fail("dynamic create_group_mask result must be b32 mask"); + } FailureOr allMask = createAllTrueMask(loc, maskType, rewriter); - if (failed(allMask)) + if (failed(allMask)) { return fail("failed to create dynamic create_group_mask all mask"); + } Value chunkBase = createI32Constant(loc, chunk * *lanesPerPart, rewriter); Value indexInPart = @@ -2993,8 +3298,9 @@ FailureOr> materializeDynamicGroupMaskForType( FailureOr laneInGroup = createPowerOfTwoRemainder( loc, logicalLane, groupSize, *allMask, rewriter); - if (failed(laneInGroup)) + if (failed(laneInGroup)) { return fail("failed to compute dynamic create_group_mask lane index"); + } Value predicate = rewriter @@ -3008,8 +3314,9 @@ FailureOr> materializeDynamicGroupMaskForType( for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { FailureOr padding = isPaddingLane(resultVMIType, part, chunk, lane); - if (failed(padding)) + if (failed(padding)) { return fail("failed to classify dynamic create_group_mask padding"); + } validLanes.push_back(*padding ? 0 : 1); hasPadding |= *padding; } @@ -3037,8 +3344,9 @@ std::optional getPrefixActiveLaneCount(ArrayRef activeLanes) { int64_t activeCount = 0; for (int8_t active : activeLanes) { if (active) { - if (seenInactive) + if (seenInactive) { return std::nullopt; + } ++activeCount; continue; } @@ -3053,13 +3361,15 @@ FailureOr materializePrefixMask(Location loc, MaskType maskType, PatternRewriter &rewriter) { std::optional pattern = getPrefixPattern(activeLanes, lanesPerPart); - if (pattern) + if (pattern) { return createPatternMask(loc, maskType, *pattern, rewriter); + } FailureOr> maskAndRemaining = createRuntimePrefixMask( loc, maskType, createI32Constant(loc, activeLanes, rewriter), rewriter); - if (failed(maskAndRemaining)) + if (failed(maskAndRemaining)) { return failure(); + } return maskAndRemaining->first; } @@ -3078,33 +3388,39 @@ FailureOr materializeConstantMaskChunk(Location loc, MaskType maskType, rewriter); FailureOr allTrue = createAllTrueMask(loc, maskType, rewriter); - if (failed(allTrue)) + if (failed(allTrue)) { return failure(); + } Value result; int64_t lane = 0; while (lane < *lanesPerPart) { - while (lane < *lanesPerPart && !activeLanes[lane]) + while (lane < *lanesPerPart && !activeLanes[lane]) { ++lane; - if (lane >= *lanesPerPart) + } + if (lane >= *lanesPerPart) { break; + } int64_t runBegin = lane; - while (lane < *lanesPerPart && activeLanes[lane]) + while (lane < *lanesPerPart && activeLanes[lane]) { ++lane; + } int64_t runEnd = lane; FailureOr prefixEnd = materializePrefixMask(loc, maskType, runEnd, *lanesPerPart, rewriter); - if (failed(prefixEnd)) + if (failed(prefixEnd)) { return failure(); + } Value runMask = *prefixEnd; if (runBegin != 0) { FailureOr prefixBegin = materializePrefixMask( loc, maskType, runBegin, *lanesPerPart, rewriter); - if (failed(prefixBegin)) + if (failed(prefixBegin)) { return failure(); + } Value notPrefixBegin = rewriter.create(loc, maskType, *prefixBegin, *allTrue) .getResult(); @@ -3122,8 +3438,9 @@ FailureOr materializeConstantMaskChunk(Location loc, MaskType maskType, .getResult(); } - if (result) + if (result) { return result; + } return materializePrefixMask(loc, maskType, 0, *lanesPerPart, rewriter); } @@ -3133,8 +3450,9 @@ FailureOr createScalarOffsetConstant(Location loc, Type type, Value createChunkOffset(Location loc, Value baseOffset, int64_t laneOffset, PatternRewriter &rewriter) { - if (laneOffset == 0) + if (laneOffset == 0) { return baseOffset; + } Value delta = rewriter.create(loc, laneOffset); return rewriter.create(loc, baseOffset, delta).getResult(); } @@ -3160,13 +3478,16 @@ LogicalResult checkContiguousFullGroupChunks( }; VMILayoutAttr layout = type.getLayoutAttr(); - if (!layout || !layout.isContiguous()) + if (!layout || !layout.isContiguous()) { return fail("group op requires contiguous VMI layout"); - if (failed(checkFullDataPhysicalChunks(type, nullptr))) + } + if (failed(checkFullDataPhysicalChunks(type, nullptr))) { return fail("group op requires full physical chunks"); + } FailureOr lanes = getDataLanesPerPart(type.getElementType()); - if (failed(lanes)) + if (failed(lanes)) { return fail("group op requires known physical lanes per part"); + } if (groupSize <= 0 || type.getElementCount() % groupSize != 0) return fail("group op requires derived group size to evenly divide lane " "count"); @@ -3185,8 +3506,9 @@ FailureOr createZeroVector(Location loc, VRegType type, FailureOr zero = createScalarOffsetConstant(loc, type.getElementType(), 0, rewriter); FailureOr mask = createAllTrueMaskForVReg(loc, type, rewriter); - if (failed(zero) || failed(mask)) + if (failed(zero) || failed(mask)) { return failure(); + } return rewriter .create(loc, type, *zero, *mask, /*position=*/nullptr) @@ -3198,11 +3520,13 @@ FailureOr createLaneRangeMask(Location loc, MaskType maskType, PatternRewriter &rewriter) { FailureOr lanesPerPart = getMaskLanesPerPart(maskType.getGranularity()); - if (failed(lanesPerPart) || begin < 0 || begin > end || end > *lanesPerPart) + if (failed(lanesPerPart) || begin < 0 || begin > end || end > *lanesPerPart) { return failure(); + } SmallVector active(*lanesPerPart, 0); - for (int64_t lane = begin; lane < end; ++lane) + for (int64_t lane = begin; lane < end; ++lane) { active[lane] = 1; + } return materializeConstantMaskChunk(loc, maskType, active, rewriter); } @@ -3218,16 +3542,19 @@ FailureOr createGroupSlotIndexVector(Location loc, VRegType indexType, FailureOr maskType = getMaskTypeForVReg(indexType, rewriter.getContext()); FailureOr allMask = createAllTrueMaskForVReg(loc, indexType, rewriter); - if (failed(baseScalar) || failed(maskType) || failed(allMask)) + if (failed(baseScalar) || failed(maskType) || failed(allMask)) { return failure(); + } Value result = rewriter .create(loc, indexType, *baseScalar, *allMask, /*position=*/nullptr) .getResult(); - if (groupSize >= lanesPerPart) + if (groupSize >= lanesPerPart) { return result; - if (lanesPerPart % groupSize != 0) + } + if (lanesPerPart % groupSize != 0) { return failure(); + } int64_t groupsPerChunk = lanesPerPart / groupSize; for (int64_t localGroup = 1; localGroup < groupsPerChunk; ++localGroup) { @@ -3237,8 +3564,9 @@ FailureOr createGroupSlotIndexVector(Location loc, VRegType indexType, FailureOr laneMask = createLaneRangeMask(loc, *maskType, localGroup * groupSize, (localGroup + 1) * groupSize, rewriter); - if (failed(groupScalar) || failed(laneMask)) + if (failed(groupScalar) || failed(laneMask)) { return failure(); + } Value splat = rewriter .create(loc, indexType, *groupScalar, *allMask, /*position=*/nullptr) @@ -3252,96 +3580,115 @@ FailureOr createGroupSlotIndexVector(Location loc, VRegType indexType, std::optional getX2MemoryDistToken(Type elementType, StringRef prefix) { unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); - if (elementBits != 8 && elementBits != 16 && elementBits != 32) + if (elementBits != 8 && elementBits != 16 && elementBits != 32) { return std::nullopt; + } return (Twine(prefix) + "_B" + Twine(elementBits)).str(); } std::optional getDenseLaneStrideLoadDistToken(VMIVRegType type) { VMILayoutAttr layout = type.getLayoutAttr(); - if (!layout || !layout.isContiguous()) + if (!layout || !layout.isContiguous()) { return std::nullopt; + } unsigned elementBits = pto::getPTOStorageElemBitWidth(type.getElementType()); if (layout.getLaneStride() == 2 && (elementBits == 8 || elementBits == 16 || elementBits == 32)) return (Twine("UNPK_B") + Twine(elementBits)).str(); - if (layout.getLaneStride() == 4 && elementBits == 8) + if (layout.getLaneStride() == 4 && elementBits == 8) { return std::string("UNPK4"); + } return std::nullopt; } std::optional getLaneStrideStoreDistToken(VMILayoutAttr layout, Type elementType) { - if (!layout || !layout.hasLaneStride()) + if (!layout || !layout.hasLaneStride()) { return std::nullopt; + } unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); - if (layout.getLaneStride() == 2 && elementBits == 8) + if (layout.getLaneStride() == 2 && elementBits == 8) { return std::string("PK_B16"); - if (layout.getLaneStride() == 2 && elementBits == 16) + } + if (layout.getLaneStride() == 2 && elementBits == 16) { return std::string("PK_B32"); - if (layout.getLaneStride() == 2 && elementBits == 32) + } + if (layout.getLaneStride() == 2 && elementBits == 32) { return std::string("PK_B64"); - if (layout.getLaneStride() == 4 && elementBits == 8) + } + if (layout.getLaneStride() == 4 && elementBits == 8) { return std::string("PK4_B32"); + } return std::nullopt; } std::optional getDenseLaneStrideStoreDistToken(VMIVRegType type) { VMILayoutAttr layout = type.getLayoutAttr(); - if (!layout || !layout.isContiguous()) + if (!layout || !layout.isContiguous()) { return std::nullopt; + } return getLaneStrideStoreDistToken(layout, type.getElementType()); } std::optional getLaneStrideStoreMaskGranularity(VMILayoutAttr layout, Type elementType) { - if (!layout || !layout.hasLaneStride()) + if (!layout || !layout.hasLaneStride()) { return std::nullopt; + } unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); - if (layout.getLaneStride() == 2 && elementBits == 8) + if (layout.getLaneStride() == 2 && elementBits == 8) { return StringRef("b16"); + } if (layout.getLaneStride() == 2 && (elementBits == 16 || elementBits == 32)) return StringRef("b32"); - if (layout.getLaneStride() == 4 && elementBits == 8) + if (layout.getLaneStride() == 4 && elementBits == 8) { return StringRef("b32"); + } return std::nullopt; } std::optional getDenseLaneStrideStoreMaskGranularity(VMIVRegType type) { VMILayoutAttr layout = type.getLayoutAttr(); - if (!layout || !layout.isContiguous()) + if (!layout || !layout.isContiguous()) { return std::nullopt; + } return getLaneStrideStoreMaskGranularity(layout, type.getElementType()); } std::optional getDenseLaneStrideMaskedStoreMaskGranularity(VMIVRegType type) { VMILayoutAttr layout = type.getLayoutAttr(); - if (!layout || !layout.isContiguous()) + if (!layout || !layout.isContiguous()) { return std::nullopt; + } unsigned elementBits = pto::getPTOStorageElemBitWidth(type.getElementType()); - if (layout.getLaneStride() == 2 && elementBits == 8) + if (layout.getLaneStride() == 2 && elementBits == 8) { return StringRef("b16"); - if (layout.getLaneStride() == 2 && elementBits == 16) + } + if (layout.getLaneStride() == 2 && elementBits == 16) { return StringRef("b32"); - if (layout.getLaneStride() == 4 && elementBits == 8) + } + if (layout.getLaneStride() == 4 && elementBits == 8) { return StringRef("b32"); + } return std::nullopt; } std::optional getPointStoreDistToken(Type elementType) { unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); - if (elementBits != 8 && elementBits != 16 && elementBits != 32) + if (elementBits != 8 && elementBits != 16 && elementBits != 32) { return std::nullopt; + } return (Twine("1PT_B") + Twine(elementBits)).str(); } std::optional getScalarBroadcastLoadDistToken(Type elementType) { unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); - if (elementBits != 8 && elementBits != 16 && elementBits != 32) + if (elementBits != 8 && elementBits != 16 && elementBits != 32) { return std::nullopt; + } return (Twine("BRC_B") + Twine(elementBits)).str(); } @@ -3354,24 +3701,31 @@ std::optional getVPTOCmpFMode(StringRef predicate) { if (predicate == "eq" || predicate == "ne" || predicate == "lt" || predicate == "le" || predicate == "gt" || predicate == "ge") return VPTOCmpMode{predicate, std::nullopt}; - if (predicate == "oeq") + if (predicate == "oeq") { return VPTOCmpMode{StringRef("eq"), std::nullopt}; - if (predicate == "one") + } + if (predicate == "one") { return VPTOCmpMode{StringRef("ne"), std::nullopt}; - if (predicate == "olt") + } + if (predicate == "olt") { return VPTOCmpMode{StringRef("lt"), std::nullopt}; - if (predicate == "ole") + } + if (predicate == "ole") { return VPTOCmpMode{StringRef("le"), std::nullopt}; - if (predicate == "ogt") + } + if (predicate == "ogt") { return VPTOCmpMode{StringRef("gt"), std::nullopt}; - if (predicate == "oge") + } + if (predicate == "oge") { return VPTOCmpMode{StringRef("ge"), std::nullopt}; + } return std::nullopt; } std::optional getVPTOCmpIMode(StringRef predicate) { - if (predicate == "eq" || predicate == "ne") + if (predicate == "eq" || predicate == "ne") { return VPTOCmpMode{predicate, std::nullopt}; + } if (predicate == "ult") return VPTOCmpMode{ StringRef("lt"), IntegerType::SignednessSemantics::Unsigned}; @@ -3401,10 +3755,12 @@ std::optional getVPTOCmpIMode(StringRef predicate) { template std::optional getVPTOCmpMode(StringRef predicate) { - if constexpr (std::is_same_v) + if constexpr (std::is_same_v) { return getVPTOCmpIMode(predicate); - else + } + else { return getVPTOCmpFMode(predicate); + } } template @@ -3412,15 +3768,17 @@ StringRef getSupportedComparePredicateMessage() { if constexpr (std::is_same_v) return "eq/ne, unsigned integer forms ult/ule/ugt/uge, and signed " "integer forms slt/sle/sgt/sge"; - else + else { return "eq/ne/lt/le/gt/ge and ordered FP forms oeq/one/olt/ole/ogt/oge"; + } } template LogicalResult checkSupportedComparePredicate(Operation *op, StringRef predicate) { - if (getVPTOCmpMode(predicate)) + if (getVPTOCmpMode(predicate)) { return success(); + } return op->emitError() << kVMIDiagUnsupportedPrefix << "compare predicate " << predicate << " cannot be lowered to pto.vcmp; supported predicates are " @@ -3477,8 +3835,9 @@ LogicalResult verifyIdentityPartForwarding(Operation *op, FailureOr getUnsignedCarrierVRegType(MLIRContext *ctx, unsigned elementBits) { - if (elementBits != 8 && elementBits != 16 && elementBits != 32) + if (elementBits != 8 && elementBits != 16 && elementBits != 32) { return failure(); + } auto elementType = IntegerType::get( ctx, elementBits, IntegerType::SignednessSemantics::Unsigned); return VRegType::get(ctx, 2048 / elementBits, elementType); @@ -3488,8 +3847,9 @@ FailureOr getSignednessCarrierVRegType(VRegType inputType, IntegerType::SignednessSemantics signedness) { auto inputElementType = dyn_cast(inputType.getElementType()); - if (!inputElementType) + if (!inputElementType) { return failure(); + } if ((signedness == IntegerType::SignednessSemantics::Signed && !inputElementType.isUnsigned()) || (signedness == IntegerType::SignednessSemantics::Unsigned && @@ -3503,22 +3863,26 @@ getSignednessCarrierVRegType(VRegType inputType, FailureOr bitcastVReg(Location loc, Value value, Type resultType, PatternRewriter &rewriter) { - if (value.getType() == resultType) + if (value.getType() == resultType) { return value; + } auto inputType = dyn_cast(value.getType()); auto outputType = dyn_cast(resultType); - if (!inputType || !outputType) + if (!inputType || !outputType) { return failure(); + } return rewriter.create(loc, outputType, value).getResult(); } FailureOr getVcaddResultType(VRegType inputType) { auto inputIntegerType = dyn_cast(inputType.getElementType()); - if (!inputIntegerType || inputIntegerType.getWidth() == 32) + if (!inputIntegerType || inputIntegerType.getWidth() == 32) { return inputType; + } unsigned inputWidth = inputIntegerType.getWidth(); - if (inputWidth != 8 && inputWidth != 16) + if (inputWidth != 8 && inputWidth != 16) { return failure(); + } auto resultElementType = IntegerType::get( inputType.getContext(), inputWidth * 2, inputIntegerType.getSignedness()); @@ -3531,8 +3895,9 @@ FailureOr unpackToNextCarrier(Location loc, Value source, PatternRewriter &rewriter) { FailureOr resultType = getUnsignedCarrierVRegType(rewriter.getContext(), sourceBits * 2); - if (failed(resultType)) + if (failed(resultType)) { return failure(); + } Value part = rewriter.create(loc, partIndex); return rewriter.create(loc, *resultType, source, part) .getResult(); @@ -3544,8 +3909,9 @@ FailureOr packToPreviousCarrier(Location loc, Value source, PatternRewriter &rewriter) { FailureOr resultType = getUnsignedCarrierVRegType(rewriter.getContext(), resultBits); - if (failed(resultType)) + if (failed(resultType)) { return failure(); + } return rewriter .create(loc, *resultType, source, rewriter.getStringAttr(part)) @@ -3576,38 +3942,44 @@ FailureOr> materializeContiguousToLaneStride( MLIRContext *ctx = rewriter.getContext(); FailureOr inputCarrier = getUnsignedCarrierVRegType(ctx, elementBits); - if (failed(inputCarrier)) + if (failed(inputCarrier)) { return failure(); + } SmallVector results; results.reserve(resultTypes.size()); for (auto [resultIndex, resultType] : llvm::enumerate(resultTypes)) { int64_t sourceIndex = resultIndex / laneStride; - if (sourceIndex >= static_cast(sourceParts.size())) + if (sourceIndex >= static_cast(sourceParts.size())) { return failure(); + } Value source = sourceParts[sourceIndex]; FailureOr current = bitcastVReg(op->getLoc(), source, *inputCarrier, rewriter); - if (failed(current)) + if (failed(current)) { return failure(); + } int64_t part = resultIndex % laneStride; FailureOr unpacked = unpackToNextCarrier(op->getLoc(), *current, elementBits, laneStride == 4 ? part / 2 : part, rewriter); - if (failed(unpacked)) + if (failed(unpacked)) { return failure(); + } current = *unpacked; if (laneStride == 4) { unpacked = unpackToNextCarrier(op->getLoc(), *current, elementBits * 2, part % 2, rewriter); - if (failed(unpacked)) + if (failed(unpacked)) { return failure(); + } current = *unpacked; } FailureOr result = bitcastVReg(op->getLoc(), *current, resultType, rewriter); - if (failed(result)) + if (failed(result)) { return failure(); + } results.push_back(*result); } return results; @@ -3638,8 +4010,9 @@ FailureOr> materializeLaneStrideToContiguous( static_cast(elementBits * static_cast(laneStride)); FailureOr sourceCarrier = getUnsignedCarrierVRegType(rewriter.getContext(), carrierBits); - if (failed(sourceCarrier)) + if (failed(sourceCarrier)) { return failure(); + } SmallVector results; results.reserve(resultTypes.size()); @@ -3652,8 +4025,9 @@ FailureOr> materializeLaneStrideToContiguous( for (Value source : sourceParts.slice(sourceBegin, sourceEnd - sourceBegin)) { FailureOr carrier = bitcastVReg(op->getLoc(), source, *sourceCarrier, rewriter); - if (failed(carrier)) + if (failed(carrier)) { return failure(); + } currentLevel.push_back(*carrier); } @@ -3665,8 +4039,9 @@ FailureOr> materializeLaneStrideToContiguous( FailureOr low = packToPreviousCarrier( op->getLoc(), currentLevel[index], currentBits / 2, "LOWER", rewriter); - if (failed(low)) + if (failed(low)) { return failure(); + } Value merged = *low; if (index + 1 < currentLevel.size()) { FailureOr high = packToPreviousCarrier( @@ -3674,8 +4049,9 @@ FailureOr> materializeLaneStrideToContiguous( "HIGHER", rewriter); FailureOr mask = createAllTrueMaskForVReg( op->getLoc(), cast((*low).getType()), rewriter); - if (failed(high) || failed(mask)) + if (failed(high) || failed(mask)) { return failure(); + } merged = rewriter .create(op->getLoc(), (*low).getType(), *low, *high, *mask) @@ -3686,12 +4062,14 @@ FailureOr> materializeLaneStrideToContiguous( currentLevel = std::move(nextLevel); currentBits /= 2; } - if (currentLevel.size() != 1) + if (currentLevel.size() != 1) { return failure(); + } FailureOr result = bitcastVReg(op->getLoc(), currentLevel.front(), resultType, rewriter); - if (failed(result)) + if (failed(result)) { return failure(); + } results.push_back(*result); } return results; @@ -3726,19 +4104,22 @@ FailureOr> materializeGroupSlotLaneStride( unsigned carrierBits = elementBits * sourceStride; FailureOr carrierType = getUnsignedCarrierVRegType(rewriter.getContext(), carrierBits); - if (failed(carrierType)) + if (failed(carrierType)) { return fail("failed to derive group-slot source carrier type"); + } FailureOr current = bitcastVReg(op->getLoc(), source, *carrierType, rewriter); - if (failed(current)) + if (failed(current)) { return fail("failed to bitcast group-slot source carrier"); + } int64_t currentStride = sourceStride; while (currentStride < resultStride) { FailureOr unpacked = unpackToNextCarrier( op->getLoc(), *current, carrierBits, /*partIndex=*/0, rewriter); - if (failed(unpacked)) + if (failed(unpacked)) { return fail("failed to unpack group-slot lane_stride carrier"); + } current = *unpacked; currentStride *= 2; carrierBits *= 2; @@ -3746,8 +4127,9 @@ FailureOr> materializeGroupSlotLaneStride( while (currentStride > resultStride) { FailureOr packed = packToPreviousCarrier( op->getLoc(), *current, carrierBits / 2, "LOWER", rewriter); - if (failed(packed)) + if (failed(packed)) { return fail("failed to pack group-slot lane_stride carrier"); + } current = *packed; currentStride /= 2; carrierBits /= 2; @@ -3755,8 +4137,9 @@ FailureOr> materializeGroupSlotLaneStride( FailureOr result = bitcastVReg(op->getLoc(), *current, resultType, rewriter); - if (failed(result)) + if (failed(result)) { return fail("failed to bitcast group-slot result carrier"); + } results.push_back(*result); } return results; @@ -3830,8 +4213,9 @@ FailureOr> materializeDataLayoutConversion( typesMatch = false; break; } - if (typesMatch) + if (typesMatch) { return SmallVector(inputs.begin(), inputs.end()); + } } } } @@ -3884,8 +4268,9 @@ FailureOr> materializeDataLayoutConversion( auto materialize = rewriter.create( op->getLoc(), lowType, highType, lhs, rhs); results.push_back(materialize.getLow()); - if (results.size() < resultTypes.size()) + if (results.size() < resultTypes.size()) { results.push_back(materialize.getHigh()); + } } } else { if (sourceParts.empty() || resultTypes.empty() || @@ -3949,8 +4334,9 @@ FailureOr> materializeDataLayoutConversion( counts.reserve(factor); size_t base = totalParts / static_cast(factor); size_t remainder = totalParts % static_cast(factor); - for (int64_t part = 0; part < factor; ++part) + for (int64_t part = 0; part < factor; ++part) { counts.push_back(base + (static_cast(part) < remainder ? 1 : 0)); + } return counts; }; auto getPartOffsets = [](ArrayRef counts) -> SmallVector { @@ -3982,8 +4368,9 @@ FailureOr> materializeDataLayoutConversion( SmallVector sourceCounts = getPartCounts(sourceParts.size(), 4); SmallVector sourceOffsets = getPartOffsets(sourceCounts); auto getSourcePart = [&](size_t part, size_t group) -> Value { - if (group < sourceCounts[part]) + if (group < sourceCounts[part]) { return sourceParts[sourceOffsets[part] + group]; + } return sourceParts.back(); }; @@ -4021,8 +4408,9 @@ FailureOr> materializeDataLayoutConversion( Value groupResults[] = {low.getLow(), low.getHigh(), high.getLow(), high.getHigh()}; for (Value result : groupResults) { - if (results.size() >= resultTypes.size()) + if (results.size() >= resultTypes.size()) { break; + } results.push_back(result); } } @@ -4081,14 +4469,18 @@ FailureOr> materializeDataLayoutConversion( op->getLoc(), chunkType, chunkType, low.getLow(), high.getLow()); auto odd = rewriter.create( op->getLoc(), chunkType, chunkType, low.getHigh(), high.getHigh()); - if (i < resultCounts[0]) + if (i < resultCounts[0]) { part0.push_back(even.getLow()); - if (i < resultCounts[1]) + } + if (i < resultCounts[1]) { part1.push_back(odd.getLow()); - if (i < resultCounts[2]) + } + if (i < resultCounts[2]) { part2.push_back(even.getHigh()); - if (i < resultCounts[3]) + } + if (i < resultCounts[3]) { part3.push_back(odd.getHigh()); + } } results.reserve(resultTypes.size()); results.append(part0); @@ -4122,8 +4514,9 @@ FailureOr> materializeDataLayoutConversion( FailureOr> dense = materializeDataLayoutConversion( op, sourceParts, resultTypes, sourceLayout, contiguous, sourceVMIElementType, rewriter); - if (failed(dense)) + if (failed(dense)) { return failure(); + } return materializeDataLayoutConversion(op, *dense, resultTypes, contiguous, resultLayout, sourceVMIElementType, rewriter); @@ -4157,8 +4550,9 @@ FailureOr> materializeEnsureLayoutConversion( } SmallVector resultTypes; - if (failed(typeConverter.convertType(resultType, resultTypes))) + if (failed(typeConverter.convertType(resultType, resultTypes))) { return failure(); + } return materializeDataLayoutConversion(op, sourceParts, resultTypes, sourceLayout, resultLayout, sourceType.getElementType(), rewriter); @@ -4168,8 +4562,9 @@ FailureOr> createPredicateDintlv(Location loc, Type lowType, Type highType, Value lhs, Value rhs, PatternRewriter &rewriter) { auto maskType = dyn_cast(lowType); - if (!maskType || highType != lowType) + if (!maskType || highType != lowType) { return failure(); + } if (maskType.isB8()) { auto op = rewriter.create(loc, lowType, highType, lhs, rhs); return std::make_pair(op.getLow(), op.getHigh()); @@ -4189,8 +4584,9 @@ FailureOr> createPredicateIntlv(Location loc, Type lowType, Type highType, Value lhs, Value rhs, PatternRewriter &rewriter) { auto maskType = dyn_cast(lowType); - if (!maskType || highType != lowType) + if (!maskType || highType != lowType) { return failure(); + } if (maskType.isB8()) { auto op = rewriter.create(loc, lowType, highType, lhs, rhs); return std::make_pair(op.getLow(), op.getHigh()); @@ -4431,8 +4827,9 @@ FailureOr> materializeMaskLayoutConversion( if (!allTrue) { FailureOr mask = createAllTrueMask( op->getLoc(), cast(lhs.getType()), rewriter); - if (failed(mask)) + if (failed(mask)) { return failure(); + } allTrue = *mask; } return rewriter.create(op->getLoc(), lhs.getType(), lhs, rhs, @@ -4443,8 +4840,9 @@ FailureOr> materializeMaskLayoutConversion( MaskType maskType) -> FailureOr { Value packed = rewriter.create(op->getLoc(), maskType, lowSource, lower); - if (!highSource) + if (!highSource) { return packed; + } Value higherPacked = rewriter.create( op->getLoc(), maskType, *highSource, higher); return mergeMasks(packed, higherPacked); @@ -4456,32 +4854,38 @@ FailureOr> materializeMaskLayoutConversion( return rewriter.notifyMatchFailure( op, "dense mask lane_stride pack requires mask result type"); size_t base = resultIndex * static_cast(laneStride); - if (base >= sourceParts.size()) + if (base >= sourceParts.size()) { break; + } std::optional source1; - if (base + 1 < sourceParts.size()) + if (base + 1 < sourceParts.size()) { source1 = sourceParts[base + 1]; + } FailureOr lowHalf = packPair(sourceParts[base], source1, maskType); - if (failed(lowHalf)) + if (failed(lowHalf)) { return failure(); + } Value current = *lowHalf; if (laneStride == 4) { current = rewriter.create(op->getLoc(), maskType, current, lower); if (base + 2 < sourceParts.size()) { std::optional source3; - if (base + 3 < sourceParts.size()) + if (base + 3 < sourceParts.size()) { source3 = sourceParts[base + 3]; + } FailureOr highHalf = packPair(sourceParts[base + 2], source3, maskType); - if (failed(highHalf)) + if (failed(highHalf)) { return failure(); + } Value higherPacked = rewriter.create( op->getLoc(), maskType, *highHalf, higher); FailureOr merged = mergeMasks(current, higherPacked); - if (failed(merged)) + if (failed(merged)) { return failure(); + } current = *merged; } } @@ -4500,12 +4904,15 @@ FailureOr> materializeMaskLayoutConversion( } int getMaskGranularityRank(StringRef granularity) { - if (granularity == "b8") + if (granularity == "b8") { return 0; - if (granularity == "b16") + } + if (granularity == "b16") { return 1; - if (granularity == "b32") + } + if (granularity == "b32") { return 2; + } return -1; } @@ -4526,15 +4933,18 @@ LogicalResult checkSupportedMaskGranularityMaterialization( VMIMaskType sourceType, VMIMaskType resultType, std::string *reason) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; - if (sourceType.getElementCount() != resultType.getElementCount()) + if (sourceType.getElementCount() != resultType.getElementCount()) { return fail("requires source and result mask lane counts to match"); - if (sourceType.getLayoutAttr() != resultType.getLayoutAttr()) + } + if (sourceType.getLayoutAttr() != resultType.getLayoutAttr()) { return fail("requires source and result mask layouts to match"); + } if (!VMIMaskType::isConcreteGranularity(sourceType.getGranularity()) || !VMIMaskType::isConcreteGranularity(resultType.getGranularity())) @@ -4543,10 +4953,12 @@ LogicalResult checkSupportedMaskGranularityMaterialization( FailureOr sourceArity = getVMIPhysicalArity(sourceType); FailureOr resultArity = getVMIPhysicalArity(resultType); - if (failed(sourceArity) || failed(resultArity)) + if (failed(sourceArity) || failed(resultArity)) { return fail("requires computable source/result physical arity"); - if (*sourceArity < 1 || *resultArity < 1) + } + if (*sourceArity < 1 || *resultArity < 1) { return fail("requires non-empty source/result physical arity"); + } return success(); } @@ -4561,8 +4973,9 @@ FailureOr> materializeAdjacentMaskGranularityConversion( int sourceRank = getMaskGranularityRank(sourceType.getGranularity()); int resultRank = getMaskGranularityRank(resultType.getGranularity()); - if (std::abs(sourceRank - resultRank) != 1) + if (std::abs(sourceRank - resultRank) != 1) { return fail("mask granularity conversion must be adjacent"); + } FailureOr sourceArity = getVMIPhysicalArity(sourceType); FailureOr factor = getVMITypeLayoutFactor(sourceType); @@ -4579,8 +4992,9 @@ FailureOr> materializeAdjacentMaskGranularityConversion( for (int64_t part = 0; part < *factor; ++part) { FailureOr sourceChunks = getVMITypeChunksInPart(sourceType, part); FailureOr resultChunks = getVMITypeChunksInPart(resultType, part); - if (failed(sourceChunks) || failed(resultChunks)) + if (failed(sourceChunks) || failed(resultChunks)) { return fail("requires computable source/result chunks per layout part"); + } if (resultRank > sourceRank) { int64_t produced = 0; @@ -4592,8 +5006,9 @@ FailureOr> materializeAdjacentMaskGranularityConversion( source, partAttr("LOWER")) .getResult()); ++produced; - if (produced >= *resultChunks) + if (produced >= *resultChunks) { break; + } results.push_back(rewriter .create(op->getLoc(), resultMaskType, source, partAttr("HIGHER")) @@ -4624,8 +5039,9 @@ FailureOr> materializeAdjacentMaskGranularityConversion( if (!allTrue) { FailureOr mask = createAllTrueMask(op->getLoc(), resultMaskType, rewriter); - if (failed(mask)) + if (failed(mask)) { return fail("failed to create all-true mask for ppack merge"); + } allTrue = *mask; } packed = rewriter @@ -4681,8 +5097,9 @@ FailureOr> materializeMaskGranularityConversion( FailureOr> nextParts = materializeAdjacentMaskGranularityConversion(op, currentType, nextType, currentParts, rewriter); - if (failed(nextParts)) + if (failed(nextParts)) { return failure(); + } currentType = nextType; currentParts = std::move(*nextParts); } @@ -4694,28 +5111,34 @@ FailureOr> getConvertedMaskPartTypes(VMIMaskType type) { FailureOr arity = getVMIPhysicalArity(type); FailureOr physicalGranularity = getVMIMaskPhysicalGranularity(type); - if (failed(arity) || failed(physicalGranularity) || *arity < 0) + if (failed(arity) || failed(physicalGranularity) || *arity < 0) { return failure(); + } SmallVector types; types.reserve(*arity); Type partType = MaskType::get(type.getContext(), *physicalGranularity); - for (int64_t i = 0; i < *arity; ++i) + for (int64_t i = 0; i < *arity; ++i) { types.push_back(partType); + } return types; } static FailureOr getVMIMaskPhysicalCarrierLayout(VMIMaskType type) { VMILayoutAttr layout = type.getLayoutAttr(); - if (!layout) + if (!layout) { return failure(); + } MLIRContext *ctx = type.getContext(); - if (layout.isContiguous()) + if (layout.isContiguous()) { return VMILayoutAttr::getContiguous(ctx); - if (layout.isDeinterleaved()) + } + if (layout.isDeinterleaved()) { return VMILayoutAttr::getDeinterleaved(ctx, layout.getFactor()); - if (layout.isBlockDeinterleaved()) + } + if (layout.isBlockDeinterleaved()) { return VMILayoutAttr::getBlockDeinterleaved(ctx, layout.getFactor()); + } if (layout.isGroupSlots()) return VMILayoutAttr::getGroupSlots(ctx, layout.getNumGroups(), layout.getSlots()); @@ -4728,8 +5151,9 @@ getVMIMaskPhysicalCarrierType(VMIMaskType type) { getVMIMaskPhysicalGranularity(type); FailureOr physicalLayout = getVMIMaskPhysicalCarrierLayout(type); - if (failed(physicalGranularity) || failed(physicalLayout)) + if (failed(physicalGranularity) || failed(physicalLayout)) { return failure(); + } return VMIMaskType::get(type.getContext(), type.getElementCount(), *physicalGranularity, *physicalLayout); } @@ -4743,8 +5167,9 @@ static bool isElementDeinterleavedLayout(VMILayoutAttr layout, FailureOr createAllFalseMaskLike(Location loc, Value value, PatternRewriter &rewriter) { auto maskType = dyn_cast(value.getType()); - if (!maskType) + if (!maskType) { return failure(); + } return createPrefixMask(loc, maskType, "PAT_ALLF", rewriter); } @@ -4773,11 +5198,13 @@ FailureOr> materializeStagingDeintToContiguousMaskLayout( FailureOr> materialized = createPredicateIntlv( op->getLoc(), nextType(0), nextType(1), sourceParts[i], sourceParts[groups + i], rewriter); - if (failed(materialized)) + if (failed(materialized)) { return fail("unsupported predicate intlv staging mask type"); + } results.push_back(materialized->first); - if (results.size() < resultTypes.size()) + if (results.size() < resultTypes.size()) { results.push_back(materialized->second); + } continue; } @@ -4791,26 +5218,32 @@ FailureOr> materializeStagingDeintToContiguousMaskLayout( FailureOr> odd = createPredicateIntlv(op->getLoc(), nextType(0), nextType(1), p1, p3, rewriter); - if (failed(even) || failed(odd)) + if (failed(even) || failed(odd)) { return fail("unsupported predicate intlv staging mask type"); + } FailureOr> low = createPredicateIntlv( op->getLoc(), nextType(0), nextType(1), even->first, odd->first, rewriter); FailureOr> high = createPredicateIntlv( op->getLoc(), nextType(2), nextType(3), even->second, odd->second, rewriter); - if (failed(low) || failed(high)) + if (failed(low) || failed(high)) { return fail("unsupported predicate intlv staging mask type"); + } results.push_back(low->first); - if (results.size() < resultTypes.size()) + if (results.size() < resultTypes.size()) { results.push_back(low->second); - if (results.size() < resultTypes.size()) + } + if (results.size() < resultTypes.size()) { results.push_back(high->first); - if (results.size() < resultTypes.size()) + } + if (results.size() < resultTypes.size()) { results.push_back(high->second); + } } - if (results.size() != resultTypes.size()) + if (results.size() != resultTypes.size()) { return fail("staging deinterleaved mask layout result arity mismatch"); + } return results; } @@ -4826,17 +5259,20 @@ FailureOr> materializeStagingContiguousToDeintMaskLayout( return fail("staging contiguous mask layout requires grouped result parts"); int64_t groups = resultTypes.size() / factor; - if (sourceParts.size() > static_cast(groups * factor)) + if (sourceParts.size() > static_cast(groups * factor)) { return fail("staging contiguous mask layout has too many source parts"); + } SmallVector, 4> parts(factor); - for (int64_t part = 0; part < factor; ++part) + for (int64_t part = 0; part < factor; ++part) { parts[part].reserve(groups); + } for (int64_t i = 0; i < groups; ++i) { size_t sourceBase = static_cast(i * factor); - if (sourceBase >= sourceParts.size()) + if (sourceBase >= sourceParts.size()) { return fail("staging contiguous mask layout ran out of source parts"); + } SmallVector sources; sources.reserve(factor); @@ -4849,8 +5285,9 @@ FailureOr> materializeStagingContiguousToDeintMaskLayout( FailureOr zero = createAllFalseMaskLike(op->getLoc(), sourceParts[sourceBase], rewriter); - if (failed(zero)) + if (failed(zero)) { return fail("failed to create all-false staging mask"); + } sources.push_back(*zero); } @@ -4859,8 +5296,9 @@ FailureOr> materializeStagingContiguousToDeintMaskLayout( createPredicateDintlv(op->getLoc(), resultTypes[i], resultTypes[groups + i], sources[0], sources[1], rewriter); - if (failed(materialized)) + if (failed(materialized)) { return fail("unsupported predicate dintlv staging mask type"); + } parts[0].push_back(materialized->first); parts[1].push_back(materialized->second); continue; @@ -4872,16 +5310,18 @@ FailureOr> materializeStagingContiguousToDeintMaskLayout( FailureOr> high = createPredicateDintlv( op->getLoc(), resultTypes[2 * groups + i], resultTypes[3 * groups + i], sources[2], sources[3], rewriter); - if (failed(low) || failed(high)) + if (failed(low) || failed(high)) { return fail("unsupported predicate dintlv staging mask type"); + } FailureOr> even = createPredicateDintlv( op->getLoc(), resultTypes[i], resultTypes[2 * groups + i], low->first, high->first, rewriter); FailureOr> odd = createPredicateDintlv( op->getLoc(), resultTypes[groups + i], resultTypes[3 * groups + i], low->second, high->second, rewriter); - if (failed(even) || failed(odd)) + if (failed(even) || failed(odd)) { return fail("unsupported predicate dintlv staging mask type"); + } parts[0].push_back(even->first); parts[1].push_back(odd->first); parts[2].push_back(even->second); @@ -4891,8 +5331,9 @@ FailureOr> materializeStagingContiguousToDeintMaskLayout( SmallVector results; results.reserve(resultTypes.size()); for (int64_t part = 0; part < factor; ++part) { - if (parts[part].size() != static_cast(groups)) + if (parts[part].size() != static_cast(groups)) { return fail("staging contiguous mask layout result arity mismatch"); + } results.append(parts[part]); } return results; @@ -4912,14 +5353,16 @@ materializeMaskGranularityCastLayoutConversionViaContiguous( sourceType.getGranularity(), contiguous); FailureOr> contiguousTypes = getConvertedMaskPartTypes(contiguousType); - if (failed(contiguousTypes)) + if (failed(contiguousTypes)) { return failure(); + } FailureOr> contiguousParts = materializeMaskGranularityCastLayoutConversion( op, sourceType, contiguousType, sourceParts, *contiguousTypes, rewriter); - if (failed(contiguousParts)) + if (failed(contiguousParts)) { return failure(); + } return materializeMaskGranularityCastLayoutConversion( op, contiguousType, resultType, *contiguousParts, resultTypes, rewriter); } @@ -4934,8 +5377,9 @@ FailureOr> materializeMaskGranularityCastLayoutConversion( VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); VMILayoutAttr resultLayout = resultType.getLayoutAttr(); - if (!sourceLayout || !resultLayout) + if (!sourceLayout || !resultLayout) { return fail("mask granularity cast layout conversion requires layouts"); + } if (sourceLayout == resultLayout) { if (failed(verifyIdentityPartForwarding(op, sourceParts, resultTypes, @@ -4946,8 +5390,9 @@ FailureOr> materializeMaskGranularityCastLayoutConversion( FailureOr> layoutParts = materializeMaskLayoutConversion( op, sourceParts, resultTypes, sourceLayout, resultLayout, rewriter); - if (succeeded(layoutParts)) + if (succeeded(layoutParts)) { return layoutParts; + } bool sourceC = sourceLayout.isContiguous() && sourceLayout.getLaneStride() == 1; bool resultC = resultLayout.isContiguous() && resultLayout.getLaneStride() == 1; @@ -4979,15 +5424,17 @@ FailureOr> materializeMaskGranularityCastConversion( return failure(); }; - if (sourceType.getElementCount() != resultType.getElementCount()) + if (sourceType.getElementCount() != resultType.getElementCount()) { return fail("requires source and result mask lane counts to match"); + } FailureOr physicalSourceType = getVMIMaskPhysicalCarrierType(sourceType); FailureOr physicalResultType = getVMIMaskPhysicalCarrierType(resultType); - if (failed(physicalSourceType) || failed(physicalResultType)) + if (failed(physicalSourceType) || failed(physicalResultType)) { return fail("requires source/result mask physical carrier types"); + } if (*physicalSourceType == *physicalResultType) { if (failed(verifyIdentityPartForwarding(op, sourceParts, resultTypes, @@ -5009,8 +5456,9 @@ FailureOr> materializeMaskGranularityCastConversion( materializeMaskGranularityConversion(op, *physicalSourceType, granularityType, sourceParts, rewriter); - if (failed(granularityParts)) + if (failed(granularityParts)) { return failure(); + } return materializeMaskGranularityCastLayoutConversion( op, granularityType, *physicalResultType, *granularityParts, resultTypes, rewriter); @@ -5028,8 +5476,9 @@ struct OneToNVMIEnsureLayoutOpPattern FailureOr> results = materializeEnsureLayoutConversion( op, adaptor.getSource(), sourceType, resultType, *this->getTypeConverter(), rewriter); - if (failed(results)) + if (failed(results)) { return failure(); + } replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); return success(); } @@ -5062,13 +5511,15 @@ struct OneToNVMIEnsureMaskLayoutOpPattern ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); FailureOr> results = materializeMaskLayoutConversion( op, sourceParts, resultTypes, sourceLayout, resultLayout, rewriter); - if (failed(results)) + if (failed(results)) { return failure(); + } replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); return success(); } @@ -5099,15 +5550,17 @@ struct OneToNVMIEnsureMaskGranularityOpPattern ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); FailureOr> results = materializeMaskGranularityCastConversion( op, sourceType, resultType, sourceParts, resultTypes, rewriter); - if (failed(results)) + if (failed(results)) { return failure(); + } if (results->size() != resultTypes.size()) return rewriter.notifyMatchFailure( op, "mask granularity cast result arity mismatch"); @@ -5140,17 +5593,19 @@ struct OneToNVMIBroadcastOpPattern : OpConversionPattern { getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); SmallVector results; results.reserve(resultTypes.size()); for (Type resultType : resultTypes) { auto vregType = dyn_cast(resultType); - if (!vregType) + if (!vregType) { return rewriter.notifyMatchFailure(op, "broadcast result must be vreg"); + } FailureOr mask = createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); if (failed(mask)) @@ -5189,22 +5644,26 @@ FailureOr createScalarOffsetConstant(Location loc, Type type, FailureOr createIotaChunkBase(Location loc, Value base, int64_t laneOffset, StringRef order, PatternRewriter &rewriter) { - if (laneOffset == 0) + if (laneOffset == 0) { return base; + } FailureOr offset = createScalarOffsetConstant(loc, base.getType(), laneOffset, rewriter); - if (failed(offset)) + if (failed(offset)) { return failure(); + } if (isa(base.getType())) { - if (order == "DESC") + if (order == "DESC") { return rewriter.create(loc, base, *offset).getResult(); + } return rewriter.create(loc, base, *offset).getResult(); } if (isa(base.getType())) { - if (order == "DESC") + if (order == "DESC") { return rewriter.create(loc, base, *offset).getResult(); + } return rewriter.create(loc, base, *offset).getResult(); } @@ -5221,8 +5680,9 @@ FailureOr createIotaContiguousChunk(Location loc, Type resultType, StringRef order = orderAttr ? orderAttr.getValue() : StringRef("ASC"); FailureOr chunkBase = createIotaChunkBase(loc, base, laneOffset, order, rewriter); - if (failed(chunkBase)) + if (failed(chunkBase)) { return failure(); + } return rewriter.create(loc, resultType, *chunkBase, orderAttr) .getResult(); } @@ -5245,17 +5705,20 @@ FailureOr createSubVLGroupPeriodicChunk(Location loc, Type resultType, StringAttr orderAttr, PatternRewriter &rewriter) { auto vregType = dyn_cast(resultType); - if (!vregType) + if (!vregType) { return failure(); + } int64_t lanesPerPart = vregType.getElementCount(); - if (groupSize <= 0 || lanesPerPart % groupSize != 0) + if (groupSize <= 0 || lanesPerPart % groupSize != 0) { return failure(); + } FailureOr allMask = createAllTrueMaskForVReg(loc, vregType, rewriter); - if (failed(allMask)) + if (failed(allMask)) { return failure(); + } // group_size==1: dst[i] = base for every lane — broadcast, not a ramp pack. if (groupSize == 1) { @@ -5279,8 +5742,9 @@ FailureOr createSubVLGroupPeriodicChunk(Location loc, Type resultType, createScalarOffsetConstant(loc, base.getType(), 0, rewriter); FailureOr maskScalar = createScalarOffsetConstant( loc, base.getType(), groupSize - 1, rewriter); - if (failed(zeroScalar) || failed(maskScalar)) + if (failed(zeroScalar) || failed(maskScalar)) { return failure(); + } Value laneIds = rewriter.create(loc, resultType, *zeroScalar, StringAttr{}) @@ -5313,8 +5777,9 @@ FailureOr createSubVLGroupPeriodicChunk(Location loc, Type resultType, getMaskTypeForVReg(vregType, rewriter.getContext()); FailureOr zeroScalar = createScalarOffsetConstant(loc, base.getType(), 0, rewriter); - if (failed(full) || failed(maskType) || failed(zeroScalar)) + if (failed(full) || failed(maskType) || failed(zeroScalar)) { return failure(); + } Value result = rewriter .create(loc, resultType, *zeroScalar, *allMask, @@ -5326,8 +5791,9 @@ FailureOr createSubVLGroupPeriodicChunk(Location loc, Type resultType, int64_t delta = localGroup * groupSize; FailureOr offsetScalar = createScalarOffsetConstant(loc, base.getType(), delta, rewriter); - if (failed(offsetScalar)) + if (failed(offsetScalar)) { return failure(); + } // ASC continuous is base+i; lane (g*S+j) holds base+g*S+j, want base+j // → subtract g*S. DESC continuous is base-i; want base-j → add g*S. if (order == "DESC") { @@ -5353,8 +5819,9 @@ FailureOr createSubVLGroupPeriodicChunk(Location loc, Type resultType, FailureOr laneMask = createLaneRangeMask(loc, *maskType, localGroup * groupSize, (localGroup + 1) * groupSize, rewriter); - if (failed(laneMask)) + if (failed(laneMask)) { return failure(); + } result = rewriter .create(loc, resultType, adjusted, result, *laneMask) .getResult(); @@ -5369,16 +5836,18 @@ FailureOr createIotaDeinterleavedChunk(Location loc, Type resultType, StringAttr orderAttr, PatternRewriter &rewriter) { auto vregType = dyn_cast(resultType); - if (!vregType) + if (!vregType) { return failure(); + } FailureOr mask = createAllTrueMaskForVReg(loc, vregType, rewriter); FailureOr zero = createScalarOffsetConstant(loc, base.getType(), 0, rewriter); FailureOr factorScalar = createScalarOffsetConstant(loc, base.getType(), factor, rewriter); - if (failed(mask) || failed(zero) || failed(factorScalar)) + if (failed(mask) || failed(zero) || failed(factorScalar)) { return failure(); + } Value local = rewriter.create(loc, resultType, *zero, StringAttr{}).getResult(); @@ -5390,8 +5859,9 @@ FailureOr createIotaDeinterleavedChunk(Location loc, Type resultType, int64_t partOffset = part + factor * chunk * lanesPerPart; FailureOr biasedBase = createIotaChunkBase(loc, base, partOffset, order, rewriter); - if (failed(biasedBase)) + if (failed(biasedBase)) { return failure(); + } if (order == "DESC") { Value baseVector = rewriter @@ -5417,8 +5887,9 @@ struct OneToNVMIIotaOpPattern : OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto resultVMIType = cast(op.getResult().getType()); VMILayoutAttr layout = resultVMIType.getLayoutAttr(); - if (!layout) + if (!layout) { return rewriter.notifyMatchFailure(op, "iota requires assigned layout"); + } FailureOr lanesPerPart = getDataLanesPerPart(resultVMIType.getElementType()); @@ -5428,16 +5899,18 @@ struct OneToNVMIIotaOpPattern : OpConversionPattern { FailureOr base = getSingleValue( op, adaptor.getBase(), "iota base must convert to one value", rewriter); - if (failed(base)) + if (failed(base)) { return failure(); + } FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); SmallVector results; @@ -5487,8 +5960,9 @@ struct OneToNVMIIotaOpPattern : OpConversionPattern { // laneOffset = (p * physVL) % S and are shared by that key. llvm::DenseMap, Value> sharedChunks; for (auto [index, resultType] : llvm::enumerate(resultTypes)) { - if (!isa(resultType)) + if (!isa(resultType)) { return rewriter.notifyMatchFailure(op, "iota result must be vreg"); + } int64_t laneOffset = 0; if (groupSizeMultipleOfPhys) @@ -5522,8 +5996,9 @@ struct OneToNVMIIotaOpPattern : OpConversionPattern { if (layout.isContiguous()) { for (auto [index, resultType] : llvm::enumerate(resultTypes)) { - if (!isa(resultType)) + if (!isa(resultType)) { return rewriter.notifyMatchFailure(op, "iota result must be vreg"); + } FailureOr result = createIotaContiguousChunk( op.getLoc(), resultType, *base, static_cast(index) * *lanesPerPart, op.getOrderAttr(), @@ -5572,8 +6047,9 @@ struct OneToNVMIConstantOpPattern : OpConversionPattern { return rewriter.notifyMatchFailure( op, "only splat dense data constants are supported"); auto splatAttr = dyn_cast(denseAttr.getSplatValue()); - if (!splatAttr) + if (!splatAttr) { return rewriter.notifyMatchFailure(op, "splat constant must be typed"); + } // arith.constant only accepts signless integer types, whereas VMI vregs may // carry signed/unsigned element types (e.g. ui16). Remap an unsigned/signed @@ -5590,15 +6066,17 @@ struct OneToNVMIConstantOpPattern : OpConversionPattern { rewriter.create(op.getLoc(), splatAttr).getResult(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); SmallVector results; results.reserve(resultTypes.size()); for (Type resultType : resultTypes) { auto vregType = dyn_cast(resultType); - if (!vregType) + if (!vregType) { return rewriter.notifyMatchFailure(op, "constant result must be vreg"); + } FailureOr mask = createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); if (failed(mask)) @@ -5625,14 +6103,16 @@ struct OneToNVMIConstantMaskOpPattern ConversionPatternRewriter &rewriter) const override { FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); std::string reason; FailureOr> materializations = computeConstantMaskMaterialization(op, &reason); - if (failed(materializations)) + if (failed(materializations)) { return rewriter.notifyMatchFailure(op, Twine("constant_mask ") + reason); + } SmallVector results; results.reserve(resultTypes.size()); @@ -5690,16 +6170,18 @@ struct OneToNVMICreateMaskOpPattern FailureOr active = getSingleValue( op, adaptor.getActiveLanes(), "create_mask active_lanes must convert to one value", rewriter); - if (failed(active)) + if (failed(active)) { return failure(); + } FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); int64_t factor = layout.isDenseSplit() ? layout.getFactor() : 1; @@ -5743,18 +6225,21 @@ struct OneToNVMICreateMaskOpPattern op, "create_mask active_lanes must be an integer constant"); int64_t activeLanes = activeAttr.getInt(); - if (activeLanes < 0) + if (activeLanes < 0) { activeLanes = 0; - if (activeLanes > resultVMIType.getElementCount()) + } + if (activeLanes > resultVMIType.getElementCount()) { activeLanes = resultVMIType.getElementCount(); + } FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); int64_t factor = layout.isDenseSplit() ? layout.getFactor() : 1; @@ -5771,19 +6256,22 @@ struct OneToNVMICreateMaskOpPattern if (failed(padding)) return rewriter.notifyMatchFailure( op, "failed to map create_mask physical padding lane"); - if (*padding) + if (*padding) { continue; + } anyLane = true; FailureOr logicalLane = mapPhysicalLaneToLogical(resultVMIType, part, chunk, lane); if (failed(logicalLane)) return rewriter.notifyMatchFailure( op, "failed to map create_mask physical lane"); - if (*logicalLane < activeLanes) + if (*logicalLane < activeLanes) { ++activeInChunk; + } } - if (!anyLane) + if (!anyLane) { break; + } if (results.size() >= resultTypes.size()) return rewriter.notifyMatchFailure( @@ -5834,8 +6322,9 @@ struct OneToNVMICreateGroupMaskOpPattern ConversionPatternRewriter &rewriter) const override { FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); auto resultVMIType = cast(op.getResult().getType()); VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); @@ -5882,13 +6371,15 @@ struct OneToNVMICreateGroupMaskOpPattern "create_group_mask active_elems_per_group must convert to one " "value", rewriter); - if (failed(active)) + if (failed(active)) { return failure(); + } FailureOr> dynamicParts = materializeDynamicGroupMaskForType(op, *active, contiguousType, resultTypes, rewriter); - if (failed(dynamicParts)) + if (failed(dynamicParts)) { return failure(); + } contiguousParts = std::move(*dynamicParts); } @@ -5898,8 +6389,9 @@ struct OneToNVMICreateGroupMaskOpPattern FailureOr> results = materializeMaskLayoutConversion( op, contiguousParts, resultTypes, contiguousLayout, resultLayout, rewriter); - if (failed(results)) + if (failed(results)) { return failure(); + } replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); return success(); } @@ -5911,8 +6403,9 @@ struct OneToNVMICreateGroupMaskOpPattern op, adaptor.getActiveElemsPerGroup(), "create_group_mask active_elems_per_group must convert to one value", rewriter); - if (failed(active)) + if (failed(active)) { return failure(); + } VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); if (resultLayout && resultLayout.isDeinterleaved()) { @@ -5924,13 +6417,15 @@ struct OneToNVMICreateGroupMaskOpPattern FailureOr> contiguousParts = materializeDynamicGroupMaskForType(op, *active, contiguousType, resultTypes, rewriter); - if (failed(contiguousParts)) + if (failed(contiguousParts)) { return failure(); + } FailureOr> results = materializeMaskLayoutConversion( op, *contiguousParts, resultTypes, contiguousLayout, resultLayout, rewriter); - if (failed(results)) + if (failed(results)) { return failure(); + } replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); return success(); @@ -5939,8 +6434,9 @@ struct OneToNVMICreateGroupMaskOpPattern FailureOr> results = materializeDynamicGroupMaskForType(op, *active, resultVMIType, resultTypes, rewriter); - if (failed(results)) + if (failed(results)) { return failure(); + } replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); return success(); } @@ -5992,12 +6488,14 @@ struct OneToNVMILoadOpPattern : OpConversionPattern { FailureOr offset = getSingleValue(op, adaptor.getOffset(), "load offset must convert to one value", rewriter); - if (failed(source) || failed(offset)) + if (failed(source) || failed(offset)) { return failure(); + } FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); if (std::optional dist = @@ -6006,8 +6504,9 @@ struct OneToNVMILoadOpPattern : OpConversionPattern { results.reserve(resultTypes.size()); int64_t semanticOffset = 0; for (auto [index, resultType] : llvm::enumerate(resultTypes)) { - if (!isa(resultType)) + if (!isa(resultType)) { return rewriter.notifyMatchFailure(op, "load result must be vreg"); + } Value chunkOffset = createChunkOffset(op.getLoc(), *offset, semanticOffset, rewriter); results.push_back(rewriter @@ -6029,8 +6528,9 @@ struct OneToNVMILoadOpPattern : OpConversionPattern { FailureOr lanesPerPart = verifyFullOrSafeReadVRegChunks( op, resultVMIType, op.getSource().getType(), *offset, rewriter); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return failure(); + } VMILayoutAttr contiguousLayout = VMILayoutAttr::getContiguous(rewriter.getContext()); @@ -6147,8 +6647,9 @@ struct OneToNVMILoadOpPattern : OpConversionPattern { contiguousParts.reserve(contiguousTypes.size()); for (auto [index, resultType] : llvm::enumerate(contiguousTypes)) { auto vregType = dyn_cast(resultType); - if (!vregType) + if (!vregType) { return rewriter.notifyMatchFailure(op, "load result must be vreg"); + } Value chunkOffset = createChunkOffset(op.getLoc(), *offset, index * *lanesPerPart, rewriter); contiguousParts.push_back(rewriter @@ -6163,8 +6664,9 @@ struct OneToNVMILoadOpPattern : OpConversionPattern { op, contiguousParts, resultTypes, contiguousLayout, resultVMIType.getLayoutAttr(), resultVMIType.getElementType(), rewriter); - if (failed(results)) + if (failed(results)) { return failure(); + } replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); return success(); @@ -6186,8 +6688,9 @@ struct OneToNVMIDeinterleaveLoadOpPattern FailureOr offset = getSingleValue( op, adaptor.getOffset(), "deinterleave_load offset must convert to one value", rewriter); - if (failed(source) || failed(offset)) + if (failed(source) || failed(offset)) { return failure(); + } FailureOr lanesPerPart = getDataLanesPerPart(lowVMIType.getElementType()); @@ -6205,15 +6708,17 @@ struct OneToNVMIDeinterleaveLoadOpPattern getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_lowTypes)) + if (failed(maybe_lowTypes)) { return failure(); + } SmallVector lowTypes = std::move(*maybe_lowTypes); FailureOr> maybe_highTypes = getConvertedResultTypes(op, 1, *this->getTypeConverter()); - if (failed(maybe_highTypes)) + if (failed(maybe_highTypes)) { return failure(); + } SmallVector highTypes = std::move(*maybe_highTypes); if (lowTypes.size() != highTypes.size()) return rewriter.notifyMatchFailure( @@ -6265,8 +6770,9 @@ struct OneToNVMIGroupLoadOpPattern : OpConversionPattern { FailureOr rowStride = getSingleValue( op, adaptor.getRowStride(), "group_load row_stride must convert to one value", rewriter); - if (failed(source) || failed(offset) || failed(rowStride)) + if (failed(source) || failed(offset) || failed(rowStride)) { return failure(); + } VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); if (resultLayout && resultLayout.isBlockDeinterleaved() && @@ -6301,9 +6807,10 @@ struct OneToNVMIGroupLoadOpPattern : OpConversionPattern { getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); int64_t factor = resultLayout.getFactor(); @@ -6389,8 +6896,9 @@ struct OneToNVMIGroupLoadOpPattern : OpConversionPattern { FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); SmallVector results; results.reserve(resultTypes.size()); @@ -6432,13 +6940,15 @@ struct OneToNVMIGroupLoadOpPattern : OpConversionPattern { getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); - if (static_cast(resultTypes.size()) != groupCount * chunksPerGroup) + if (static_cast(resultTypes.size()) != groupCount * chunksPerGroup) { return rewriter.notifyMatchFailure(op, "group_load arity mismatch"); + } SmallVector results; results.reserve(resultTypes.size()); @@ -6479,8 +6989,9 @@ static LogicalResult lowerGroupSlotLoadParts( int64_t slots = layout.getSlots(); int64_t expectedArity = ceilDivNonNegative(numGroups, slots); - if (static_cast(resultTypes.size()) != expectedArity) + if (static_cast(resultTypes.size()) != expectedArity) { return rewriter.notifyMatchFailure(op, "group_slot_load arity mismatch"); + } auto makeI16 = [&](int64_t value) -> Value { return rewriter.create(op->getLoc(), value, 16); @@ -6621,8 +7132,9 @@ static LogicalResult lowerGroupBroadcastParts( Operation *op, ValueRange sourceParts, VMIVRegType sourceVMIType, VMIVRegType resultVMIType, TypeRange resultTypes, int64_t numGroups, ConversionPatternRewriter &rewriter, SmallVectorImpl &results) { - if (sourceParts.empty() || resultTypes.empty()) + if (sourceParts.empty() || resultTypes.empty()) { return rewriter.notifyMatchFailure(op, "group_broadcast arity mismatch"); + } std::string layoutReason; VMILayoutSupport supports; @@ -6710,14 +7222,16 @@ static LogicalResult lowerGroupBroadcastParts( auto getSelector = [&](int64_t baseSlot) -> FailureOr { int64_t baseIndex = baseSlot * sourceLaneStride; auto cached = selectorByBaseIndex.find(baseIndex); - if (cached != selectorByBaseIndex.end()) + if (cached != selectorByBaseIndex.end()) { return cached->second; + } if (selectorKind == SelectorKind::Constant) { FailureOr baseScalar = createScalarOffsetConstant( op->getLoc(), indexScalarType, baseIndex, rewriter); - if (failed(baseScalar)) + if (failed(baseScalar)) { return failure(); + } Value selector = rewriter .create(op->getLoc(), indexType, *baseScalar, *allMask, @@ -6730,8 +7244,9 @@ static LogicalResult lowerGroupBroadcastParts( if (!sharedRamp) { FailureOr zero = createScalarOffsetConstant( op->getLoc(), indexScalarType, 0, rewriter); - if (failed(zero)) + if (failed(zero)) { return failure(); + } sharedRamp = rewriter.create(op->getLoc(), indexType, *zero, StringAttr{}) .getResult(); @@ -6756,8 +7271,9 @@ static LogicalResult lowerGroupBroadcastParts( if (baseIndex != 0) { FailureOr baseScalar = createScalarOffsetConstant( op->getLoc(), indexScalarType, baseIndex, rewriter); - if (failed(baseScalar)) + if (failed(baseScalar)) { return failure(); + } selector = rewriter .create(op->getLoc(), indexType, selector, *baseScalar, *allMask) @@ -6826,8 +7342,9 @@ static LogicalResult lowerGroupBroadcastParts( if (failed(padding)) return rewriter.notifyMatchFailure( op, "group_broadcast failed to map result padding lanes"); - if (*padding) + if (*padding) { continue; + } FailureOr logical = mapPhysicalLaneToLogical(resultVMIType, part, chunk, lane); if (failed(logical)) @@ -6864,8 +7381,9 @@ static LogicalResult lowerGroupBroadcastParts( for (int64_t chunkIndex : llvm::drop_begin(activeSourceChunks)) { SmallVector laneMaskBits(fact->lanesPerPart, 0); for (auto [lane, laneSourceChunk] : llvm::enumerate(laneSourceChunks)) - if (laneSourceChunk == chunkIndex) + if (laneSourceChunk == chunkIndex) { laneMaskBits[lane] = 1; + } FailureOr laneMask = materializeConstantMaskChunk( op->getLoc(), *resultMaskType, laneMaskBits, rewriter); if (failed(laneMask)) @@ -6891,8 +7409,9 @@ static LogicalResult lowerGroupBroadcastParts( if (failed(padding)) return rewriter.notifyMatchFailure( op, "group_broadcast failed to map result padding lanes"); - if (*padding) + if (*padding) { continue; + } FailureOr logical = mapPhysicalLaneToLogical(resultVMIType, part, chunk, lane); if (failed(logical)) @@ -6900,8 +7419,9 @@ static LogicalResult lowerGroupBroadcastParts( op, "group_broadcast failed to map a result lane"); int64_t actualGroup = *logical / fact->groupSize; int64_t expectedGroup = firstGroup; - if (selectorKind != SelectorKind::Constant) + if (selectorKind != SelectorKind::Constant) { expectedGroup += lane / selectorPeriod; + } if (actualGroup != expectedGroup || actualGroup / sourceSlots != sourceChunk) return rewriter.notifyMatchFailure( @@ -6960,16 +7480,18 @@ struct OneToNVMIGroupSlotLoadOpPattern op, adaptor.getSourceGroupStride(), "group_slot_load source_group_stride must convert to one value", rewriter); - if (failed(source) || failed(offset) || failed(sourceGroupStride)) + if (failed(source) || failed(offset) || failed(sourceGroupStride)) { return failure(); + } FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); int64_t numGroups = op.getNumGroupsAttr().getInt(); @@ -6998,20 +7520,23 @@ struct OneToNVMIMaskedLoadOpPattern FailureOr offset = getSingleValue( op, adaptor.getOffset(), "masked_load offset must convert to one value", rewriter); - if (failed(source) || failed(offset)) + if (failed(source) || failed(offset)) { return failure(); + } FailureOr lanesPerPart = verifyFullOrSafeReadVRegChunks( op, resultVMIType, (*source).getType(), *offset, rewriter); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return failure(); + } ValueRange maskParts = adaptor.getMask(); ValueRange passthruParts = adaptor.getPassthru(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (maskParts.size() != passthruParts.size() || passthruParts.size() != resultTypes.size()) @@ -7056,16 +7581,18 @@ struct OneToNVMIGatherOpPattern : OpConversionPattern { FailureOr source = getSingleValue(op, adaptor.getSource(), "gather source must convert to one value", rewriter); - if (failed(source)) + if (failed(source)) { return failure(); + } ValueRange indicesParts = adaptor.getIndices(); ValueRange maskParts = adaptor.getMask(); ValueRange passthruParts = adaptor.getPassthru(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (indicesParts.size() != maskParts.size() || indicesParts.size() != passthruParts.size() || @@ -7117,23 +7644,26 @@ struct OneToNVMIExpandLoadOpPattern FailureOr offset = getSingleValue( op, adaptor.getOffset(), "expand_load offset must convert to one value", rewriter); - if (failed(source) || failed(offset)) + if (failed(source) || failed(offset)) { return failure(); + } FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (isStaticAllActiveMask(op.getMask(), resultVMIType.getElementCount())) { FailureOr lanesPerPart = verifyFullOrSafeReadVRegChunks( op, resultVMIType, (*source).getType(), *offset, rewriter); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return failure(); + } SmallVector results; results.reserve(resultTypes.size()); @@ -7230,8 +7760,9 @@ struct OneToNVMIStoreOpPattern : OpConversionPattern { FailureOr offset = getSingleValue(op, adaptor.getOffset(), "store offset must convert to one value", rewriter); - if (failed(destination) || failed(offset)) + if (failed(destination) || failed(offset)) { return failure(); + } ValueRange valueParts = adaptor.getValue(); if (std::optional dist = @@ -7244,15 +7775,17 @@ struct OneToNVMIStoreOpPattern : OpConversionPattern { int64_t semanticOffset = 0; for (auto [index, value] : llvm::enumerate(valueParts)) { auto vregType = dyn_cast(value.getType()); - if (!vregType) + if (!vregType) { return rewriter.notifyMatchFailure(op, "store value must be vreg"); + } FailureOr activeLanes = getActiveDataLanesInPhysicalChunk(valueVMIType, index); if (failed(activeLanes)) return rewriter.notifyMatchFailure( op, "failed to compute lane_stride store active lanes"); - if (*activeLanes == 0) + if (*activeLanes == 0) { continue; + } auto maskType = MaskType::get(rewriter.getContext(), *maskGranularity); FailureOr mask = createPrefixMaskForActiveLanes( op.getLoc(), maskType, *activeLanes, rewriter); @@ -7282,8 +7815,9 @@ struct OneToNVMIStoreOpPattern : OpConversionPattern { SmallVector contiguousTypes = std::move(*maybeContiguousTypes); SmallVector valuePartTypes; valuePartTypes.reserve(valueParts.size()); - for (Value value : valueParts) + for (Value value : valueParts) { valuePartTypes.push_back(value.getType()); + } FailureOr noWiderThanContiguous = hasNoWiderFootprintThanContiguous(valuePartTypes, contiguousTypes); if (failed(noWiderThanContiguous)) @@ -7307,8 +7841,9 @@ struct OneToNVMIStoreOpPattern : OpConversionPattern { return rewriter.notifyMatchFailure( op, "vstsx2 requires matching low/high value types"); auto vregType = dyn_cast(low.getType()); - if (!vregType) + if (!vregType) { return rewriter.notifyMatchFailure(op, "store value must be vreg"); + } FailureOr mask = createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); if (failed(mask)) @@ -7328,21 +7863,24 @@ struct OneToNVMIStoreOpPattern : OpConversionPattern { FailureOr> storeParts = materializeDataLayoutConversion( op, valueParts, contiguousTypes, valueVMIType.getLayoutAttr(), contiguousLayout, valueVMIType.getElementType(), rewriter); - if (failed(storeParts)) + if (failed(storeParts)) { return failure(); + } for (auto [index, value] : llvm::enumerate(*storeParts)) { auto vregType = dyn_cast(value.getType()); - if (!vregType) + if (!vregType) { return rewriter.notifyMatchFailure(op, "store value must be vreg"); + } if (!fullPhysicalChunks) { FailureOr activeLanes = getContiguousActiveDataLanes(valueVMIType, index); if (failed(activeLanes)) return rewriter.notifyMatchFailure( op, "failed to compute store active lanes"); - if (*activeLanes == 0) + if (*activeLanes == 0) { continue; + } } FailureOr mask = fullPhysicalChunks @@ -7391,8 +7929,9 @@ struct OneToNVMIInterleaveStoreOpPattern FailureOr offset = getSingleValue( op, adaptor.getOffset(), "interleave_store offset must convert to one value", rewriter); - if (failed(destination) || failed(offset)) + if (failed(destination) || failed(offset)) { return failure(); + } ValueRange lowParts = adaptor.getLow(); ValueRange highParts = adaptor.getHigh(); @@ -7447,8 +7986,9 @@ struct OneToNVMIGroupStoreOpPattern FailureOr rowStride = getSingleValue( op, adaptor.getRowStride(), "group_store row_stride must convert to one value", rewriter); - if (failed(destination) || failed(offset) || failed(rowStride)) + if (failed(destination) || failed(offset) || failed(rowStride)) { return failure(); + } unsigned elementBits = pto::getPTOStorageElemBitWidth(valueVMIType.getElementType()); @@ -7583,9 +8123,9 @@ struct OneToNVMIGroupStoreOpPattern if (static_cast(valueParts.size()) != layout.getNumGroups()) return rewriter.notifyMatchFailure( op, "slots=1 group_store arity mismatch"); - unsigned elementBits = + unsigned slots1ElementBits = pto::getPTOStorageElemBitWidth(valueVMIType.getElementType()); - if (elementBits == 0 || 256 % elementBits != 0) + if (slots1ElementBits == 0 || 256 % slots1ElementBits != 0) return rewriter.notifyMatchFailure( op, "slots=1 group_store requires supported element width"); std::optional constantRowStride = @@ -7783,12 +8323,14 @@ struct OneToNVMIGroupStoreOpPattern Value merged = *zero; for (int64_t localPart = 0; localPart < 4; ++localPart) { int64_t partIndex = blockStart / 8 + localPart; - if (partIndex >= static_cast(valueParts.size())) + if (partIndex >= static_cast(valueParts.size())) { break; + } int64_t remainingGroups = numGroups - partIndex * 8; int64_t activeGroups = std::min(8, remainingGroups); - if (activeGroups <= 0) + if (activeGroups <= 0) { break; + } Value selected = rewriter .create(op.getLoc(), firstVRegType, @@ -8000,8 +8542,9 @@ struct OneToNVMIGroupStoreOpPattern return failure(); ValueRange valueParts = adaptor.getValue(); - if (static_cast(valueParts.size()) != groupCount * chunksPerGroup) + if (static_cast(valueParts.size()) != groupCount * chunksPerGroup) { return rewriter.notifyMatchFailure(op, "group_store arity mismatch"); + } for (auto [index, value] : llvm::enumerate(valueParts)) { auto vregType = dyn_cast(value.getType()); @@ -8048,8 +8591,9 @@ struct OneToNVMIMaskedStoreOpPattern FailureOr offset = getSingleValue( op, adaptor.getOffset(), "masked_store offset must convert to one value", rewriter); - if (failed(destination) || failed(offset)) + if (failed(destination) || failed(offset)) { return failure(); + } ValueRange valueParts = adaptor.getValue(); ValueRange maskParts = adaptor.getMask(); @@ -8079,8 +8623,9 @@ struct OneToNVMIMaskedStoreOpPattern if (failed(activeLanes)) return rewriter.notifyMatchFailure( op, "failed to compute lane_stride masked_store active lanes"); - if (*activeLanes == 0) + if (*activeLanes == 0) { continue; + } FailureOr storeMask = createDenseLaneStrideStorePredicate( op.getLoc(), valueVMIType, index, mask, *maskGranularity, rewriter); @@ -8103,24 +8648,28 @@ struct OneToNVMIMaskedStoreOpPattern SmallVector contiguousValueTypes; contiguousValueTypes.reserve(valueParts.size()); - for (Value value : valueParts) + for (Value value : valueParts) { contiguousValueTypes.push_back(value.getType()); + } FailureOr> storeParts = materializeDataLayoutConversion( op, valueParts, contiguousValueTypes, valueVMIType.getLayoutAttr(), VMILayoutAttr::getContiguous(rewriter.getContext()), valueVMIType.getElementType(), rewriter); - if (failed(storeParts)) + if (failed(storeParts)) { return failure(); + } SmallVector contiguousMaskTypes; contiguousMaskTypes.reserve(maskParts.size()); - for (Value mask : maskParts) + for (Value mask : maskParts) { contiguousMaskTypes.push_back(mask.getType()); + } FailureOr> storeMasks = materializeMaskLayoutConversion( op, maskParts, contiguousMaskTypes, maskVMIType.getLayoutAttr(), VMILayoutAttr::getContiguous(rewriter.getContext()), rewriter); - if (failed(storeMasks)) + if (failed(storeMasks)) { return failure(); + } if (storeParts->size() != storeMasks->size()) return rewriter.notifyMatchFailure( @@ -8138,8 +8687,9 @@ struct OneToNVMIMaskedStoreOpPattern if (failed(activeLanes)) return rewriter.notifyMatchFailure( op, "failed to compute masked_store active lanes"); - if (*activeLanes == 0) + if (*activeLanes == 0) { continue; + } FailureOr storeMask = createMaskedStorePredicate( op.getLoc(), valueVMIType, index, mask, vregType, rewriter); if (failed(storeMask)) @@ -8176,8 +8726,9 @@ struct OneToNVMIGroupBroadcastLoadOpPattern op, adaptor.getSourceGroupStride(), "group_broadcast_load source_group_stride must convert to one value", rewriter); - if (failed(source) || failed(offset) || failed(sourceGroupStride)) + if (failed(source) || failed(offset) || failed(sourceGroupStride)) { return failure(); + } VMILayoutSupport supports; std::string supportReason; @@ -8190,8 +8741,9 @@ struct OneToNVMIGroupBroadcastLoadOpPattern FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); FailureOr directFact = @@ -8199,12 +8751,15 @@ struct OneToNVMIGroupBroadcastLoadOpPattern auto getBRCDist = [&]() -> std::optional { unsigned elementBits = pto::getPTOStorageElemBitWidth(resultVMIType.getElementType()); - if (elementBits == 8) + if (elementBits == 8) { return StringRef("BRC_B8"); - if (elementBits == 16) + } + if (elementBits == 16) { return StringRef("BRC_B16"); - if (elementBits == 32) + } + if (elementBits == 32) { return StringRef("BRC_B32"); + } return std::nullopt; }; @@ -8411,8 +8966,9 @@ struct OneToNVMIStrideLoadOpPattern ValueRange maskParts = adaptor.getMask(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (resultTypes.size() != 1 || maskParts.size() != 1) return rewriter.notifyMatchFailure( @@ -8492,8 +9048,9 @@ struct OneToNVMIScatterOpPattern : OpConversionPattern { FailureOr destination = getSingleValue( op, adaptor.getDestination(), "scatter destination must convert to one value", rewriter); - if (failed(destination)) + if (failed(destination)) { return failure(); + } ValueRange valueParts = adaptor.getValue(); ValueRange indicesParts = adaptor.getIndices(); @@ -8529,8 +9086,9 @@ struct OneToNVMIBinaryOpPattern : OpConversionPattern { ValueRange rhsParts = adaptor.getRhs(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (lhsParts.size() != rhsParts.size() || lhsParts.size() != resultTypes.size()) @@ -8580,8 +9138,9 @@ struct OneToNVMIVecScalarOpPattern : OpConversionPattern { ValueRange maskParts = adaptor.getMask(); FailureOr> maybeResultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(scalar) || failed(maybeResultTypes)) + if (failed(scalar) || failed(maybeResultTypes)) { return failure(); + } Value scalarValue = *scalar; SmallVector resultTypes = std::move(*maybeResultTypes); if (sourceParts.empty() || sourceParts.size() != maskParts.size() || @@ -8740,8 +9299,9 @@ struct OneToNVMIVmullOpPattern : OpConversionPattern { getConvertedResultTypes(op, 0, *this->getTypeConverter()); FailureOr> maybeHighTypes = getConvertedResultTypes(op, 1, *this->getTypeConverter()); - if (failed(maybeLowTypes) || failed(maybeHighTypes)) + if (failed(maybeLowTypes) || failed(maybeHighTypes)) { return failure(); + } SmallVector lowTypes = std::move(*maybeLowTypes); SmallVector highTypes = std::move(*maybeHighTypes); @@ -8806,8 +9366,9 @@ struct OneToNVMIInterleaveOpPattern : OpConversionPattern { getConvertedResultTypes(op, 0, *this->getTypeConverter()); FailureOr> maybeHighTypes = getConvertedResultTypes(op, 1, *this->getTypeConverter()); - if (failed(maybeLowTypes) || failed(maybeHighTypes)) + if (failed(maybeLowTypes) || failed(maybeHighTypes)) { return failure(); + } SmallVector lowTypes = std::move(*maybeLowTypes); SmallVector highTypes = std::move(*maybeHighTypes); if (lhsParts.size() != rhsParts.size() || @@ -8841,10 +9402,12 @@ struct OneToNVMIInterleaveOpPattern : OpConversionPattern { return layout && layout.isContiguous() && layout.getLaneStride() == 1; }; auto getElementDeintFactor = [](VMILayoutAttr layout) -> int64_t { - if (layout && layout.isContiguous() && layout.getLaneStride() == 1) + if (layout && layout.isContiguous() && layout.getLaneStride() == 1) { return 1; - if (layout && layout.isDeinterleaved() && layout.getLaneStride() == 1) + } + if (layout && layout.isDeinterleaved() && layout.getLaneStride() == 1) { return layout.getFactor(); + } return 0; }; @@ -9020,8 +9583,9 @@ struct OneToNVMIFmaOpPattern : OpConversionPattern { ValueRange accParts = adaptor.getAcc(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (lhsParts.size() != rhsParts.size() || lhsParts.size() != accParts.size() || @@ -9068,8 +9632,9 @@ struct OneToNVMIVexpdifOpPattern : OpConversionPattern { ValueRange maskParts = adaptor.getMask(); FailureOr> maybeResultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybeResultTypes)) + if (failed(maybeResultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybeResultTypes); if (xParts.size() != maxParts.size() || xParts.size() != maskParts.size() || @@ -9111,11 +9676,13 @@ struct OneToNVMIUnaryOpPattern : OpConversionPattern { ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); - if (sourceParts.size() != resultTypes.size()) + if (sourceParts.size() != resultTypes.size()) { return rewriter.notifyMatchFailure(op, "physical unary arity mismatch"); + } SmallVector results; results.reserve(resultTypes.size()); @@ -9152,8 +9719,9 @@ struct OneToNVMIMaskBinaryOpPattern : OpConversionPattern { ValueRange rhsParts = adaptor.getRhs(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (lhsParts.size() != rhsParts.size() || lhsParts.size() != resultTypes.size()) @@ -9196,8 +9764,9 @@ struct OneToNVMIMaskUnaryOpPattern : OpConversionPattern { ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (sourceParts.size() != resultTypes.size()) return rewriter.notifyMatchFailure(op, @@ -9247,8 +9816,9 @@ struct OneToNVMICmpOpPattern : OpConversionPattern { ValueRange rhsParts = adaptor.getRhs(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (lhsParts.size() != rhsParts.size() || lhsParts.size() != resultTypes.size()) @@ -9307,8 +9877,9 @@ struct OneToNVMISelectOpPattern : OpConversionPattern { ValueRange falseParts = adaptor.getFalseValue(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (maskParts.size() != trueParts.size() || trueParts.size() != falseParts.size() || @@ -9344,8 +9915,9 @@ struct OneToNVMIVselrOpPattern : OpConversionPattern { ValueRange indexParts = adaptor.getIndex(); FailureOr> maybeResultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybeResultTypes)) + if (failed(maybeResultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybeResultTypes); if (sourceParts.size() != 1 || indexParts.size() != 1 || @@ -9386,8 +9958,9 @@ struct OneToNVMIActivePrefixIndexOpPattern ValueRange maskParts = adaptor.getMask(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (maskParts.size() != 1 || resultTypes.size() != 1) return rewriter.notifyMatchFailure( @@ -9437,8 +10010,9 @@ struct OneToNVMICompressOpPattern : OpConversionPattern { ValueRange maskParts = adaptor.getMask(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (sourceParts.size() != 1 || maskParts.size() != 1 || resultTypes.size() != 1) @@ -9475,8 +10049,9 @@ struct OneToNVMICompressStoreOpPattern FailureOr offset = getSingleValue( op, adaptor.getOffset(), "compress_store offset must convert to one value", rewriter); - if (failed(destination) || failed(offset)) + if (failed(destination) || failed(offset)) { return failure(); + } ValueRange valueParts = adaptor.getValue(); ValueRange maskParts = adaptor.getMask(); @@ -9522,8 +10097,9 @@ struct OneToNVMIReduceAddIOpPattern ValueRange maskParts = adaptor.getMask(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (sourceParts.empty() || sourceParts.size() != maskParts.size() || resultTypes.size() != 1) @@ -9609,8 +10185,9 @@ struct OneToNVMIReduceAddFOpPattern ValueRange maskParts = adaptor.getMask(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (sourceParts.empty() || sourceParts.size() != maskParts.size() || resultTypes.size() != 1) @@ -9701,8 +10278,9 @@ classifyGroupReduceLoweringPlan(VMIVRegType sourceType, VMIMaskType maskType, FailureOr fact = supports.getGroupReduceLayoutFactForLayouts( sourceType, maskType, resultType, numGroups, reason); - if (failed(fact)) + if (failed(fact)) { return failure(); + } switch (fact->blockClass) { case VMIGroupBlockClass::QuarterBlock: @@ -9737,8 +10315,9 @@ struct OneToNVMIGroupReduceOpPattern : OpConversionPattern { ValueRange maskParts = adaptor.getMask(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); VMILayoutSupport supports; @@ -10204,8 +10783,9 @@ struct OneToNVMIGroupReduceOpPattern : OpConversionPattern { if (rowLocalSlots1Result) { results[destChunk] = *finalResult; } else { - for (int64_t chunk = 0; chunk < chunksPerGroup; ++chunk) + for (int64_t chunk = 0; chunk < chunksPerGroup; ++chunk) { results[destChunk + chunk] = *finalResult; + } } } @@ -10216,8 +10796,9 @@ struct OneToNVMIGroupReduceOpPattern : OpConversionPattern { private: FailureOr getRowResultType(VRegType sourceType, VRegType resultType) const { - if constexpr (std::is_same_v) + if constexpr (std::is_same_v) { return getVcaddResultType(sourceType); + } return resultType; } @@ -10266,8 +10847,9 @@ struct OneToNVMIGroupBroadcastOpPattern ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); SmallVector results; if (failed(lowerGroupBroadcastParts( @@ -10306,8 +10888,9 @@ lowerVMIHistogramToVPTO(VMIOp op, op, "expected matching source/mask chunks"); auto partType = dyn_cast(accParts[0].getType()); - if (!partType) + if (!partType) { return rewriter.notifyMatchFailure(op, "expected ui16 acc parts"); + } if (halfCount == 2 && accParts[1].getType() != partType) return rewriter.notifyMatchFailure(op, "expected matching ui16 acc parts"); @@ -10315,14 +10898,16 @@ lowerVMIHistogramToVPTO(VMIOp op, auto sourceType = cast(op.getSource().getType()); FailureOr lanesPerPart = getDataLanesPerPart(sourceType.getElementType()); - if (failed(lanesPerPart)) + if (failed(lanesPerPart)) { return rewriter.notifyMatchFailure(op, "failed to compute source lanes"); + } Location loc = op.getLoc(); SmallVector binConsts; binConsts.push_back(createI32Constant(loc, 0, rewriter)); - if (halfCount == 2) + if (halfCount == 2) { binConsts.push_back(createI32Constant(loc, 1, rewriter)); + } SmallVector halves(accParts.begin(), accParts.end()); @@ -10330,8 +10915,9 @@ lowerVMIHistogramToVPTO(VMIOp op, Value source = sourceParts[index]; Value userMask = maskParts[index]; auto maskType = dyn_cast(userMask.getType()); - if (!maskType || !maskType.isB8()) + if (!maskType || !maskType.isB8()) { return rewriter.notifyMatchFailure(op, "expected b8 source mask"); + } Value chunkMask = userMask; int64_t firstLane = int64_t(index) * *lanesPerPart; @@ -10397,8 +10983,9 @@ struct OneToNVMIReduceMinMaxOpPattern : OpConversionPattern { ValueRange maskParts = adaptor.getMask(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (sourceParts.empty() || sourceParts.size() != maskParts.size() || resultTypes.size() != 1) @@ -10482,16 +11069,18 @@ struct OneToNVMIExtFOpPattern : OpConversionPattern { ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (sourceParts.empty()) return rewriter.notifyMatchFailure( op, "extf requires at least one physical source chunk"); auto sourceType = dyn_cast(sourceParts.front().getType()); - if (!sourceType) + if (!sourceType) { return rewriter.notifyMatchFailure(op, "expected physical extf source"); + } for (Value sourcePart : sourceParts) { auto currentSourceType = dyn_cast(sourcePart.getType()); if (!currentSourceType || currentSourceType != sourceType) @@ -10560,8 +11149,9 @@ struct OneToNVMIExtFOpPattern : OpConversionPattern { FailureOr mask = createAllTrueMaskForVReg(op.getLoc(), sourceType, rewriter); - if (failed(mask)) + if (failed(mask)) { return rewriter.notifyMatchFailure(op, "failed to build extf seed mask"); + } SmallVector results; results.reserve(resultTypes.size()); @@ -10594,8 +11184,9 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); @@ -10655,12 +11246,14 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { return success(); } - if (resultTypes.empty()) + if (resultTypes.empty()) { return rewriter.notifyMatchFailure(op, "truncf requires result chunks"); + } auto sourceType0 = dyn_cast(sourceParts.front().getType()); - if (!sourceType0 || !isa(sourceType0.getElementType())) + if (!sourceType0 || !isa(sourceType0.getElementType())) { return rewriter.notifyMatchFailure(op, "unsupported physical truncf source type"); + } unsigned sourceBits = pto::getPTOStorageElemBitWidth(sourceType0.getElementType()); if (sourceBits != 32 && sourceBits != 16) return rewriter.notifyMatchFailure( @@ -10698,8 +11291,9 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { sourceParts.size() == resultTypes.size()) { FailureOr sourceMask = createAllTrueMaskForVReg(op.getLoc(), sourceType0, rewriter); - if (failed(sourceMask)) + if (failed(sourceMask)) { return rewriter.notifyMatchFailure(op, "failed to build truncf masks"); + } StringAttr rnd = rewriter.getStringAttr( getTruncFRoundMode(op, resultVRegTypes.front().getElementType())); StringAttr sat = op->getAttrOfType("saturate"); @@ -10723,8 +11317,9 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { resultLayout.getLaneStride() != 1 && sourceParts.size() == resultTypes.size()) { StringRef part; - if (resultBits == 16 && resultLayout.getLaneStride() == 2) + if (resultBits == 16 && resultLayout.getLaneStride() == 2) { part = "EVEN"; // 32→16 + } else if (resultBits == 8 && resultLayout.getLaneStride() == 4) part = "P0"; // 32→8 (f8/hif8) else if (resultBits == 8 && resultLayout.getLaneStride() == 2) @@ -10735,8 +11330,9 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { FailureOr sourceMask = createAllTrueMaskForVReg(op.getLoc(), sourceType0, rewriter); - if (failed(sourceMask)) + if (failed(sourceMask)) { return rewriter.notifyMatchFailure(op, "failed to build truncf masks"); + } StringAttr rnd = rewriter.getStringAttr( getTruncFRoundMode(op, resultVRegTypes.front().getElementType())); @@ -10786,8 +11382,9 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { FailureOr sourceMask = createAllTrueMaskForVReg(op.getLoc(), sourceType0, rewriter); - if (failed(sourceMask)) + if (failed(sourceMask)) { return rewriter.notifyMatchFailure(op, "failed to build truncf masks"); + } StringAttr rnd = rewriter.getStringAttr( getTruncFRoundMode(op, resultVRegTypes.front().getElementType())); @@ -10842,8 +11439,9 @@ struct OneToNVMIExtIOpPattern : OpConversionPattern { ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (sourceParts.empty()) return rewriter.notifyMatchFailure( @@ -11059,8 +11657,9 @@ struct OneToNVMITruncIOpPattern : OpConversionPattern { ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); @@ -11292,8 +11891,9 @@ struct OneToNVMITruncIOpPattern : OpConversionPattern { StringAttr part = rewriter.getStringAttr(factor == 2 ? "EVEN" : "P0"); FailureOr sourceMask = createAllTrueMaskForVReg(op.getLoc(), sourceType0, rewriter); - if (failed(sourceMask)) + if (failed(sourceMask)) { return rewriter.notifyMatchFailure(op, "failed to build trunci masks"); + } SmallVector results; results.reserve(resultTypes.size()); @@ -11330,8 +11930,9 @@ struct OneToNVMITruncIOpPattern : OpConversionPattern { createAllTrueMaskForVReg(op.getLoc(), sourceType0, rewriter); FailureOr resultMask = createAllTrueMaskForVReg(op.getLoc(), resultType0, rewriter); - if (failed(sourceMask) || failed(resultMask)) + if (failed(sourceMask) || failed(resultMask)) { return rewriter.notifyMatchFailure(op, "failed to build trunci masks"); + } SmallVector results; results.reserve(resultTypes.size()); @@ -11375,8 +11976,9 @@ struct OneToNVMIFPToSIOpPattern : OpConversionPattern { ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); Type srcElem = sourceVMIType.getElementType(); @@ -11587,8 +12189,9 @@ struct OneToNVMIFPToUIOpPattern : OpConversionPattern { ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); Type srcElem = sourceVMIType.getElementType(); @@ -11770,8 +12373,9 @@ struct OneToNVMISIToFPOpPattern : OpConversionPattern { ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); if (sourceParts.size() != resultTypes.size()) return rewriter.notifyMatchFailure( @@ -11793,8 +12397,9 @@ struct OneToNVMISIToFPOpPattern : OpConversionPattern { FailureOr mask = createAllTrueMaskForVReg(op.getLoc(), sourceType, rewriter); - if (failed(mask)) + if (failed(mask)) { return rewriter.notifyMatchFailure(op, "failed to build sitofp mask"); + } results.push_back(rewriter .create(op.getLoc(), resultVRegType, sourcePart, *mask, rnd, @@ -11816,11 +12421,13 @@ struct OneToNVMIBitcastOpPattern : OpConversionPattern { ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); - if (sourceParts.size() != resultTypes.size()) + if (sourceParts.size() != resultTypes.size()) { return rewriter.notifyMatchFailure(op, "physical bitcast arity mismatch"); + } SmallVector results; results.reserve(resultTypes.size()); @@ -11871,15 +12478,17 @@ struct OneToNVMIChannelSplitOpPattern FailureOr> maybe_resultTypes = getConvertedResultTypes(op, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); FailureOr> results = materializeDataLayoutConversion(op, adaptor.getSource(), resultTypes, sourceLayout, channelLayout, sourceType.getElementType(), rewriter); - if (failed(results)) + if (failed(results)) { return failure(); + } replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); return success(); @@ -11918,15 +12527,17 @@ struct OneToNVMIChannelMergeOpPattern FailureOr> maybeResultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybeResultTypes)) + if (failed(maybeResultTypes)) { return failure(); + } FailureOr> results = materializeDataLayoutConversion( op, flattenOneToNOperands(adaptor.getOperands()), *maybeResultTypes, channelLayout, resultLayout, resultType.getElementType(), rewriter); - if (failed(results)) + if (failed(results)) { return failure(); + } replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); return success(); @@ -11942,8 +12553,9 @@ struct OneToNVMIShuffleOpPattern : OpConversionPattern { ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); - if (failed(maybe_resultTypes)) + if (failed(maybe_resultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybe_resultTypes); std::string reason; FailureOr> sourceFlatIndices = @@ -12007,8 +12619,9 @@ struct OneToNVMIShuffleOpPattern : OpConversionPattern { return rewriter.notifyMatchFailure(op, Twine("shuffle vselr ") + vselrReason); - if (vselrPlans->size() != resultTypes.size()) + if (vselrPlans->size() != resultTypes.size()) { return rewriter.notifyMatchFailure(op, "shuffle vselr arity mismatch"); + } SmallVector results; results.reserve(resultTypes.size()); @@ -12064,8 +12677,9 @@ Block *convertBranchDestBlock(Block *block, ConversionPatternRewriter &rewriter, const TypeConverter &typeConverter, llvm::DenseMap &converted) { auto [it, inserted] = converted.try_emplace(block, nullptr); - if (!inserted) + if (!inserted) { return it->second; + } TypeConverter::SignatureConversion argMapping(block->getNumArguments()); if (failed(typeConverter.convertSignatureArgs(block->getArgumentTypes(), @@ -12170,8 +12784,9 @@ struct OneToNCFSwitchOpPattern : OpConversionPattern { llvm::zip(op.getCaseOperands(), adaptor.getCaseOperands())) changed |= !isIdentityOneToNValueMapping(originalOperands, convertedOperands); - if (!changed) + if (!changed) { return failure(); + } ValueRange flag = adaptor.getFlag(); if (flag.size() != 1) @@ -12185,10 +12800,12 @@ struct OneToNCFSwitchOpPattern : OpConversionPattern { caseOperandStorage.reserve(op.getCaseOperandSegments().size()); caseOperands.reserve(op.getCaseOperandSegments().size()); - for (ArrayRef convertedOperands : adaptor.getCaseOperands()) + for (ArrayRef convertedOperands : adaptor.getCaseOperands()) { caseOperandStorage.push_back(flattenOneToNOperands(convertedOperands)); - for (SmallVector &operands : caseOperandStorage) + } + for (SmallVector &operands : caseOperandStorage) { caseOperands.push_back(operands); + } rewriter.replaceOpWithNewOp( op, flag.front(), defaultDest, defaultOperands, op.getCaseValuesAttr(), @@ -12207,11 +12824,13 @@ struct OneToNSCFExecuteRegionOpPattern ConversionPatternRewriter &rewriter) const override { FailureOr> maybeResultTypes = getConvertedResultTypes(op, *this->getTypeConverter()); - if (failed(maybeResultTypes)) + if (failed(maybeResultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybeResultTypes); - if (resultTypes == op->getResultTypes()) + if (resultTypes == op->getResultTypes()) { return failure(); + } auto newOp = rewriter.create(op.getLoc(), resultTypes); @@ -12239,11 +12858,13 @@ struct OneToNSCFIndexSwitchOpPattern FailureOr> maybeResultTypes = getConvertedResultTypes(op, *this->getTypeConverter()); - if (failed(maybeResultTypes)) + if (failed(maybeResultTypes)) { return failure(); + } SmallVector resultTypes = std::move(*maybeResultTypes); - if (resultTypes == op->getResultTypes()) + if (resultTypes == op->getResultTypes()) { return failure(); + } auto newOp = rewriter.create( op.getLoc(), resultTypes, arg.front(), op.getCases(), op.getNumCases()); @@ -12394,48 +13015,54 @@ LogicalResult verifyNoResidualVMIIR(ModuleOp module) { LogicalResult checkSupportedExtFShape(VMIExtFOp op, std::string *reason = nullptr) { VMILayoutSupport supports; - if (failed(supports.getExtFSupport(op, reason))) + if (failed(supports.getExtFSupport(op, reason))) { return failure(); + } return success(); } LogicalResult checkSupportedTruncFShape(VMITruncFOp op, std::string *reason = nullptr) { VMILayoutSupport supports; - if (failed(supports.getTruncFSupport(op, reason))) + if (failed(supports.getTruncFSupport(op, reason))) { return failure(); + } return success(); } LogicalResult checkSupportedExtSIShape(VMIExtSIOp op, std::string *reason = nullptr) { VMILayoutSupport supports; - if (failed(supports.getExtSISupport(op, reason))) + if (failed(supports.getExtSISupport(op, reason))) { return failure(); + } return success(); } LogicalResult checkSupportedExtUIShape(VMIExtUIOp op, std::string *reason = nullptr) { VMILayoutSupport supports; - if (failed(supports.getExtUISupport(op, reason))) + if (failed(supports.getExtUISupport(op, reason))) { return failure(); + } return success(); } LogicalResult checkSupportedTruncIShape(VMITruncIOp op, std::string *reason = nullptr) { VMILayoutSupport supports; - if (failed(supports.getTruncISupport(op, reason))) + if (failed(supports.getTruncISupport(op, reason))) { return failure(); + } return success(); } LogicalResult checkSupportedFPToSIShape(VMIFPToSIOp op, std::string *reason = nullptr) { auto fail = [&](const Twine &message) { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -12443,22 +13070,25 @@ LogicalResult checkSupportedFPToSIShape(VMIFPToSIOp op, auto resultType = cast(op.getResult().getType()); VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); VMILayoutAttr resultLayout = resultType.getLayoutAttr(); - if (!sourceLayout || !resultLayout) + if (!sourceLayout || !resultLayout) { return fail("requires assigned source/result layouts"); + } Type srcElem = sourceType.getElementType(); Type dstElem = resultType.getElementType(); auto contract = lookupVMIFpToSiContract(srcElem, dstElem); - if (!contract) + if (!contract) { return fail("unsupported fp-to-si conversion element type pair"); + } unsigned srcBits = pto::getPTOStorageElemBitWidth(srcElem); unsigned dstBits = pto::getPTOStorageElemBitWidth(dstElem); if (srcBits == dstBits) { // Same-width (f32→s32, f16→s16): layout equality + arity equality. - if (sourceLayout != resultLayout) + if (sourceLayout != resultLayout) { return fail("same-width fp-to-si requires matching layouts"); + } FailureOr sourceArity = getVMIPhysicalArity(sourceType); FailureOr resultArity = getVMIPhysicalArity(resultType); if (failed(sourceArity) || failed(resultArity) || @@ -12470,8 +13100,9 @@ LogicalResult checkSupportedFPToSIShape(VMIFPToSIOp op, FailureOr fact = layoutSupport.getCastLayoutFactForLayouts( sourceType, resultType, sourceLayout, resultLayout, reason); - if (failed(fact)) + if (failed(fact)) { return failure(); + } } return success(); @@ -12480,8 +13111,9 @@ LogicalResult checkSupportedFPToSIShape(VMIFPToSIOp op, LogicalResult checkSupportedFPToUIShape(VMIFPToUIOp op, std::string *reason = nullptr) { auto fail = [&](const Twine &message) { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -12489,22 +13121,25 @@ LogicalResult checkSupportedFPToUIShape(VMIFPToUIOp op, auto resultType = cast(op.getResult().getType()); VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); VMILayoutAttr resultLayout = resultType.getLayoutAttr(); - if (!sourceLayout || !resultLayout) + if (!sourceLayout || !resultLayout) { return fail("requires assigned source/result layouts"); + } Type srcElem = sourceType.getElementType(); Type dstElem = resultType.getElementType(); auto contract = lookupVMIFpToUIContract(srcElem, dstElem); - if (!contract) + if (!contract) { return fail("unsupported fp-to-ui conversion element type pair"); + } unsigned srcBits = pto::getPTOStorageElemBitWidth(srcElem); unsigned dstBits = pto::getPTOStorageElemBitWidth(dstElem); if (srcBits == dstBits) { // Same-width: layout equality + arity equality. - if (sourceLayout != resultLayout) + if (sourceLayout != resultLayout) { return fail("same-width fp-to-ui requires matching layouts"); + } FailureOr sourceArity = getVMIPhysicalArity(sourceType); FailureOr resultArity = getVMIPhysicalArity(resultType); if (failed(sourceArity) || failed(resultArity) || @@ -12516,8 +13151,9 @@ LogicalResult checkSupportedFPToUIShape(VMIFPToUIOp op, FailureOr fact = layoutSupport.getCastLayoutFactForLayouts( sourceType, resultType, sourceLayout, resultLayout, reason); - if (failed(fact)) + if (failed(fact)) { return failure(); + } } return success(); @@ -12526,8 +13162,9 @@ LogicalResult checkSupportedFPToUIShape(VMIFPToUIOp op, LogicalResult checkSupportedSIToFPShape(VMISIToFPOp op, std::string *reason = nullptr) { auto fail = [&](const Twine &message) { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -12535,15 +13172,18 @@ LogicalResult checkSupportedSIToFPShape(VMISIToFPOp op, auto resultType = cast(op.getResult().getType()); VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); VMILayoutAttr resultLayout = resultType.getLayoutAttr(); - if (!sourceLayout || !resultLayout) + if (!sourceLayout || !resultLayout) { return fail("requires assigned source/result layouts"); - if (sourceLayout != resultLayout) + } + if (sourceLayout != resultLayout) { return fail("requires source/result layouts to match"); + } if (!isa(sourceType.getElementType()) || pto::getPTOStorageElemBitWidth(sourceType.getElementType()) != 32) return fail("requires 32-bit integer source element type"); - if (!resultType.getElementType().isF32()) + if (!resultType.getElementType().isF32()) { return fail("requires f32 result element type"); + } FailureOr sourceArity = getVMIPhysicalArity(sourceType); FailureOr resultArity = getVMIPhysicalArity(resultType); if (failed(sourceArity) || failed(resultArity) || @@ -12554,8 +13194,9 @@ LogicalResult checkSupportedSIToFPShape(VMISIToFPOp op, LogicalResult checkSupportedBitcastShape(VMIBitcastOp op, std::string *reason) { VMILayoutSupport supports; - if (failed(supports.getBitcastSupport(op, reason))) + if (failed(supports.getBitcastSupport(op, reason))) { return failure(); + } return success(); } @@ -12563,19 +13204,22 @@ LogicalResult checkSupportedChannelSplitShape(VMIChannelSplitOp op, std::string *reason = nullptr) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; int64_t channels = op.getNumResults(); - if (channels != 2 && channels != 4) + if (channels != 2 && channels != 4) { return fail("pto.vmi.channel_split supports only 2 or 4 channels"); + } auto sourceType = cast(op.getSource().getType()); VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); - if (!sourceLayout) + if (!sourceLayout) { return fail("requires assigned source layout"); + } auto expectedLayout = VMILayoutAttr::getDeinterleaved(op.getContext(), channels); if (!sourceLayout.isContiguous() && sourceLayout != expectedLayout) @@ -12585,8 +13229,9 @@ checkSupportedChannelSplitShape(VMIChannelSplitOp op, for (Value result : op.getResults()) { VMILayoutAttr resultLayout = cast(result.getType()).getLayoutAttr(); - if (!resultLayout || !resultLayout.isContiguous()) + if (!resultLayout || !resultLayout.isContiguous()) { return fail("requires every result layout to be contiguous"); + } } FailureOr sourceArity = getVMIPhysicalArity(sourceType); @@ -12594,14 +13239,17 @@ checkSupportedChannelSplitShape(VMIChannelSplitOp op, for (Value result : op.getResults()) { FailureOr arity = getVMIPhysicalArity(cast(result.getType())); - if (failed(arity)) + if (failed(arity)) { return fail("requires computable result physical arity"); + } resultArity += *arity; } - if (failed(sourceArity)) + if (failed(sourceArity)) { return fail("requires computable source physical arity"); - if (*sourceArity != resultArity) + } + if (*sourceArity != resultArity) { return fail("requires source and result to have the same physical arity"); + } return success(); } @@ -12610,31 +13258,36 @@ LogicalResult checkSupportedChannelMergeShape(VMIChannelMergeOp op, std::string *reason = nullptr) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; int64_t channels = op.getInputs().size(); - if (channels != 2 && channels != 4) + if (channels != 2 && channels != 4) { return fail("pto.vmi.channel_merge supports only 2 or 4 channels"); + } int64_t inputArity = 0; for (Value input : op.getInputs()) { auto inputType = cast(input.getType()); VMILayoutAttr inputLayout = inputType.getLayoutAttr(); - if (!inputLayout || !inputLayout.isContiguous()) + if (!inputLayout || !inputLayout.isContiguous()) { return fail("requires every input layout to be contiguous"); + } FailureOr arity = getVMIPhysicalArity(inputType); - if (failed(arity)) + if (failed(arity)) { return fail("requires computable input physical arity"); + } inputArity += *arity; } auto resultType = cast(op.getResult().getType()); VMILayoutAttr resultLayout = resultType.getLayoutAttr(); - if (!resultLayout) + if (!resultLayout) { return fail("requires assigned result layout"); + } auto expectedLayout = VMILayoutAttr::getDeinterleaved(op.getContext(), channels); if (!resultLayout.isContiguous() && resultLayout != expectedLayout) @@ -12642,10 +13295,12 @@ checkSupportedChannelMergeShape(VMIChannelMergeOp op, "deinterleaved channel layout"); FailureOr resultArity = getVMIPhysicalArity(resultType); - if (failed(resultArity)) + if (failed(resultArity)) { return fail("requires computable result physical arity"); - if (*resultArity != inputArity) + } + if (*resultArity != inputArity) { return fail("requires source and result to have the same physical arity"); + } return success(); } @@ -12654,8 +13309,9 @@ LogicalResult checkSupportedActivePrefixIndexShape(VMIActivePrefixIndexOp op, std::string *reason = nullptr) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -12663,10 +13319,12 @@ checkSupportedActivePrefixIndexShape(VMIActivePrefixIndexOp op, auto resultType = cast(op.getResult().getType()); VMILayoutAttr maskLayout = maskType.getLayoutAttr(); VMILayoutAttr resultLayout = resultType.getLayoutAttr(); - if (!maskLayout || !resultLayout) + if (!maskLayout || !resultLayout) { return fail("requires assigned mask and result layouts"); - if (!maskLayout.isContiguous() || !resultLayout.isContiguous()) + } + if (!maskLayout.isContiguous() || !resultLayout.isContiguous()) { return fail("requires contiguous mask and result layouts"); + } std::string resultFullReason; if (failed(checkFullDataPhysicalChunks(resultType, &resultFullReason))) @@ -12682,8 +13340,9 @@ checkSupportedActivePrefixIndexShape(VMIActivePrefixIndexOp op, FailureOr maskArity = getVMIPhysicalArity(maskType); FailureOr resultArity = getVMIPhysicalArity(resultType); - if (failed(maskArity) || failed(resultArity)) + if (failed(maskArity) || failed(resultArity)) { return fail("requires computable mask and result physical arity"); + } if (*maskArity != 1 || *resultArity != 1) return fail("requires a single physical chunk; multi-chunk prefix needs " "cross-chunk carry"); @@ -12694,8 +13353,9 @@ checkSupportedActivePrefixIndexShape(VMIActivePrefixIndexOp op, LogicalResult checkSupportedCompressShape(VMICompressOp op, std::string *reason = nullptr) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -12705,8 +13365,9 @@ LogicalResult checkSupportedCompressShape(VMICompressOp op, VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); VMILayoutAttr maskLayout = maskType.getLayoutAttr(); VMILayoutAttr resultLayout = resultType.getLayoutAttr(); - if (!sourceLayout || !maskLayout || !resultLayout) + if (!sourceLayout || !maskLayout || !resultLayout) { return fail("requires assigned source, mask, and result layouts"); + } if (!sourceLayout.isContiguous() || !maskLayout.isContiguous() || !resultLayout.isContiguous()) return fail("requires contiguous source, mask, and result layouts"); @@ -12720,8 +13381,9 @@ LogicalResult checkSupportedCompressShape(VMICompressOp op, FailureOr sourceArity = getVMIPhysicalArity(sourceType); FailureOr maskArity = getVMIPhysicalArity(maskType); FailureOr resultArity = getVMIPhysicalArity(resultType); - if (failed(sourceArity) || failed(maskArity) || failed(resultArity)) + if (failed(sourceArity) || failed(maskArity) || failed(resultArity)) { return fail("requires computable source, mask, and result physical arity"); + } if (*sourceArity != 1 || *maskArity != 1 || *resultArity != 1) return fail("requires a single physical chunk; multi-chunk compress needs " "cross-chunk compaction"); @@ -12733,8 +13395,9 @@ LogicalResult checkSupportedCompressStoreShape( VMICompressStoreOp op, std::string *reason = nullptr) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -12742,10 +13405,12 @@ LogicalResult checkSupportedCompressStoreShape( auto maskType = cast(op.getMask().getType()); VMILayoutAttr valueLayout = valueType.getLayoutAttr(); VMILayoutAttr maskLayout = maskType.getLayoutAttr(); - if (!valueLayout || !maskLayout) + if (!valueLayout || !maskLayout) { return fail("requires assigned value and mask layouts"); - if (!valueLayout.isContiguous() || !maskLayout.isContiguous()) + } + if (!valueLayout.isContiguous() || !maskLayout.isContiguous()) { return fail("requires contiguous value and mask layouts"); + } if (!isa(op.getDestination().getType())) return fail("requires !pto.ptr destination because pto.vstur is " @@ -12759,8 +13424,9 @@ LogicalResult checkSupportedCompressStoreShape( FailureOr valueArity = getVMIPhysicalArity(valueType); FailureOr maskArity = getVMIPhysicalArity(maskType); - if (failed(valueArity) || failed(maskArity)) + if (failed(valueArity) || failed(maskArity)) { return fail("requires computable value and mask physical arity"); + } if (*valueArity != 1 || *maskArity != 1) return fail("requires a single physical chunk; multi-chunk " "compress_store needs cross-chunk compaction and SQZN " @@ -12774,13 +13440,15 @@ LogicalResult checkSupportedReduceShape(OpTy op, bool requiresReassoc, std::string *reason = nullptr) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; - if (requiresReassoc && !op->hasAttr("reassoc")) + if (requiresReassoc && !op->hasAttr("reassoc")) { return fail("requires reassoc attr for pair-wise floating-point vcadd"); + } auto sourceType = cast(op.getSource().getType()); auto maskType = cast(op.getMask().getType()); @@ -12788,11 +13456,13 @@ checkSupportedReduceShape(OpTy op, bool requiresReassoc, VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); VMILayoutAttr maskLayout = maskType.getLayoutAttr(); VMILayoutAttr resultLayout = resultType.getLayoutAttr(); - if (!sourceLayout || !maskLayout || !resultLayout) + if (!sourceLayout || !maskLayout || !resultLayout) { return fail("requires assigned source, mask, and result layouts"); + } if (!sourceLayout.isContiguous() || !maskLayout.isContiguous() || - !resultLayout.isContiguous()) + !resultLayout.isContiguous()) { return fail("requires contiguous source, mask, and result layouts"); + } std::string fullChunkReason; if (failed(checkFullDataPhysicalChunks(sourceType, &fullChunkReason))) @@ -12808,8 +13478,9 @@ checkSupportedReduceShape(OpTy op, bool requiresReassoc, if (*sourceArity < 1 || *maskArity != *sourceArity) return fail("requires source and mask physical arity to match and be " "non-empty"); - if (*resultArity != 1) + if (*resultArity != 1) { return fail("requires one result physical chunk"); + } return success(); } @@ -12819,23 +13490,29 @@ LogicalResult checkSupportedGroupReduceShape(OpTy op, std::string *reason = nullptr) { VMILayoutSupport supports; if constexpr (std::is_same_v) { - if (succeeded(supports.getGroupReduceAddFSupport(op, reason))) + if (succeeded(supports.getGroupReduceAddFSupport(op, reason))) { return success(); + } } else if constexpr (std::is_same_v) { - if (succeeded(supports.getGroupReduceMaxFSupport(op, reason))) + if (succeeded(supports.getGroupReduceMaxFSupport(op, reason))) { return success(); + } } else if constexpr (std::is_same_v) { - if (succeeded(supports.getGroupReduceMaxISupport(op, reason))) + if (succeeded(supports.getGroupReduceMaxISupport(op, reason))) { return success(); + } } else if constexpr (std::is_same_v) { - if (succeeded(supports.getGroupReduceMinFSupport(op, reason))) + if (succeeded(supports.getGroupReduceMinFSupport(op, reason))) { return success(); + } } else if constexpr (std::is_same_v) { - if (succeeded(supports.getGroupReduceMinISupport(op, reason))) + if (succeeded(supports.getGroupReduceMinISupport(op, reason))) { return success(); + } } else { - if (succeeded(supports.getGroupReduceAddISupport(op, reason))) + if (succeeded(supports.getGroupReduceAddISupport(op, reason))) { return success(); + } } return failure(); } @@ -12846,31 +13523,39 @@ LogicalResult checkSupportedGroupBroadcastShape( auto sourceType = cast(op.getSource().getType()); auto resultType = cast(op.getResult().getType()); if (sourceType.getElementType() != resultType.getElementType()) { - if (reason) + if (reason) { *reason = "requires source/result element type to match"; + } return failure(); } auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); VMILayoutAttr resultLayout = resultType.getLayoutAttr(); - if (!sourceLayout || !resultLayout) + if (!sourceLayout || !resultLayout) { return fail("requires assigned source/result layouts"); + } int64_t numGroups = op.getNumGroupsAttr().getInt(); - if (numGroups <= 0) + if (numGroups <= 0) { return fail("requires positive num_groups"); - if (sourceType.getElementCount() != numGroups) + } + if (sourceType.getElementCount() != numGroups) { return fail("requires source lane count to match num_groups"); - if (resultType.getElementCount() % numGroups != 0) + } + if (resultType.getElementCount() % numGroups != 0) { return fail("requires num_groups to evenly divide result lane count"); - if (!sourceLayout.isGroupSlots() || sourceLayout.getNumGroups() != numGroups) + } + if (!sourceLayout.isGroupSlots() || sourceLayout.getNumGroups() != numGroups) { return fail("requires matching num_groups source layout"); - if (resultLayout.isGroupSlots()) + } + if (resultLayout.isGroupSlots()) { return fail("requires dense result layout"); + } if (sourceLayout.getSlots() > 0 && sourceLayout.getSlots() != 8 && sourceLayout.getSlots() != 1) @@ -12878,8 +13563,9 @@ LogicalResult checkSupportedGroupBroadcastShape( "layouts"); VMILayoutSupport supports; std::string supportReason; - if (failed(supports.getGroupBroadcastSupport(op, &supportReason))) + if (failed(supports.getGroupBroadcastSupport(op, &supportReason))) { return fail(supportReason); + } FailureOr lanesPerPart = getDataLanesPerPart(sourceType.getElementType()); @@ -12890,15 +13576,17 @@ LogicalResult checkSupportedGroupBroadcastShape( return fail("requires matching physical lanes per part"); FailureOr groupSize = getGroupSizeFromNumGroups( resultType, numGroups, reason); - if (failed(groupSize)) + if (failed(groupSize)) { return failure(); + } if (*lanesPerPart % *groupSize != 0 && *groupSize % *lanesPerPart != 0) return fail("requires derived group size to divide or be a multiple of " "physical lanes per part"); FailureOr resultFactor = getDataLayoutFactor(resultType); - if (failed(resultFactor)) + if (failed(resultFactor)) { return fail("requires known result layout factor"); + } bool laneStridedDense = resultLayout.isDense() && resultLayout.getLaneStride() > 1; if (!laneStridedDense) { @@ -12907,8 +13595,9 @@ LogicalResult checkSupportedGroupBroadcastShape( return fail(Twine("requires full result physical chunks; ") + fullChunkReason); } - if (*resultFactor == 1) + if (*resultFactor == 1) { return success(); + } FailureOr resultBlockElems = getVMILayoutBlockElems(resultType); bool blockFragmentSmallGroup = @@ -12919,8 +13608,9 @@ LogicalResult checkSupportedGroupBroadcastShape( *groupSize < *lanesPerPart && *groupSize >= *resultFactor && *groupSize % *resultFactor == 0 && *lanesPerPart % (*groupSize / *resultFactor) == 0; - if (blockFragmentSmallGroup || deinterleavedSmallGroup) + if (blockFragmentSmallGroup || deinterleavedSmallGroup) { return success(); + } int64_t logicalSpanPerResultChunk = *lanesPerPart * *resultFactor; if (*groupSize < *lanesPerPart || *groupSize % logicalSpanPerResultChunk != 0) return fail("deinterleaved result requires every physical result chunk to " @@ -12931,24 +13621,27 @@ LogicalResult checkSupportedGroupBroadcastShape( LogicalResult checkSupportedVdhistShape(VMIVdhistOp op, std::string *reason = nullptr) { VMILayoutSupport supports; - if (succeeded(supports.getVdhistSupport(op, reason))) + if (succeeded(supports.getVdhistSupport(op, reason))) { return success(); + } return failure(); } LogicalResult checkSupportedVchistShape(VMIVchistOp op, std::string *reason = nullptr) { VMILayoutSupport supports; - if (succeeded(supports.getVchistSupport(op, reason))) + if (succeeded(supports.getVchistSupport(op, reason))) { return success(); + } return failure(); } LogicalResult checkSupportedVmullShape(VMIVmullOp op, std::string *reason = nullptr) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; @@ -12962,16 +13655,19 @@ LogicalResult checkSupportedVmullShape(VMIVmullOp op, if (!elementType || elementType.getWidth() != 32 || (!elementType.isSignless() && !elementType.isUnsigned())) return fail("requires element type to be exactly i32 or ui32"); - if (aType != bType || aType != lowType || aType != highType) + if (aType != bType || aType != lowType || aType != highType) { return fail("requires identical a, b, low, and high VMI vreg types"); + } int64_t lanes = aType.getElementCount(); - if (lanes != 64 && lanes != 128 && lanes != 256) + if (lanes != 64 && lanes != 128 && lanes != 256) { return fail("requires logical lane count 64, 128, or 256"); + } VMILayoutAttr layout = aType.getLayoutAttr(); - if (!layout) + if (!layout) { return fail("requires an assigned data layout"); + } bool supportedLayout = layout.getLaneStride() == 1 && (layout.isContiguous() || @@ -12983,8 +13679,9 @@ LogicalResult checkSupportedVmullShape(VMIVmullOp op, if (maskType.getLayoutAttr() != layout) return fail("requires the mask and all four data values to share one " "layout"); - if (maskType.getGranularity() != "b32") + if (maskType.getGranularity() != "b32") { return fail("requires b32 mask granularity"); + } FailureOr aArity = getVMIPhysicalArity(aType); FailureOr bArity = getVMIPhysicalArity(bType); @@ -13081,15 +13778,17 @@ LogicalResult checkSupportedVMIAddcsShape(VMIVaddcsOp op, LogicalResult checkSupportedFmaShape(VMIFmaOp op, std::string *reason = nullptr) { auto fail = [&](const Twine &message) -> LogicalResult { - if (reason) + if (reason) { *reason = message.str(); + } return failure(); }; auto lhsType = cast(op.getLhs().getType()); FailureOr arity = getVMIPhysicalArity(lhsType); - if (failed(arity) || *arity < 1) + if (failed(arity) || *arity < 1) { return fail("requires computable non-empty physical arity"); + } return success(); } @@ -13097,8 +13796,9 @@ checkSupportedFmaShape(VMIFmaOp op, std::string *reason = nullptr) { LogicalResult checkSupportedReluShape(VMIReluOp op, std::string *reason = nullptr) { auto resultType = cast(op.getResult().getType()); - if (failed(checkSupportedMaskableVReg(resultType, reason))) + if (failed(checkSupportedMaskableVReg(resultType, reason))) { return failure(); + } return success(); } @@ -13160,8 +13860,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, auto emitMaskableUnsupported = [&](Operation *op, StringRef opName, VMIVRegType type) -> WalkResult { std::string reason; - if (succeeded(checkSupportedMaskableVReg(type, &reason))) + if (succeeded(checkSupportedMaskableVReg(type, &reason))) { return WalkResult::advance(); + } op->emitError() << kVMIDiagUnsupportedPrefix << opName @@ -13207,8 +13908,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, } if (auto hist = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedVdhistShape(hist, &reason))) + if (succeeded(checkSupportedVdhistShape(hist, &reason))) { return WalkResult::advance(); + } hist.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.vdhist requires contiguous Nx{ui8|i8} source, contiguous b8 " @@ -13218,8 +13920,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, } if (auto hist = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedVchistShape(hist, &reason))) + if (succeeded(checkSupportedVchistShape(hist, &reason))) { return WalkResult::advance(); + } hist.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.vchist requires contiguous Nx{ui8|i8} source, contiguous b8 " @@ -13248,8 +13951,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, } if (auto load = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedStrideLoadShape(load, &reason))) + if (succeeded(checkSupportedStrideLoadShape(load, &reason))) { return WalkResult::advance(); + } load.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.stride_load lowers through pto.vsldb only for one " @@ -13259,8 +13963,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, } if (auto load = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedGroupLoadShape(load, &reason))) + if (succeeded(checkSupportedGroupLoadShape(load, &reason))) { return WalkResult::advance(); + } load.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.group_load requires contiguous full result chunks, a " @@ -13306,8 +14011,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, return WalkResult::interrupt(); } std::string reason; - if (succeeded(checkSupportedMaskedLoadShape(load, &reason))) + if (succeeded(checkSupportedMaskedLoadShape(load, &reason))) { return WalkResult::advance(); + } load.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.masked_load direct lowering requires a supported memory " @@ -13318,8 +14024,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, } if (auto gather = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedGatherShape(gather, &reason))) + if (succeeded(checkSupportedGatherShape(gather, &reason))) { return WalkResult::advance(); + } gather.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.gather lowers through pto.vgather2_bc + pto.vsel only " @@ -13330,8 +14037,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, } if (auto load = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedExpandLoadShape(load, &reason))) + if (succeeded(checkSupportedExpandLoadShape(load, &reason))) { return WalkResult::advance(); + } load.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.expand_load direct lowering is currently supported for " @@ -13411,8 +14119,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, } if (auto scatter = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedScatterShape(scatter, &reason))) + if (succeeded(checkSupportedScatterShape(scatter, &reason))) { return WalkResult::advance(); + } scatter.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.scatter lowers through pto.vscatter only with a UB " @@ -13526,22 +14235,29 @@ verifySupportedVMIToVPTOOps(ModuleOp module, return emitMaskableUnsupported( op, opName, cast(vecScalar.getResult().getType())); }; - if (auto vecScalar = dyn_cast(op)) + if (auto vecScalar = dyn_cast(op)) { return verifyVecScalar(vecScalar, "pto.vmi.vadds"); - if (auto vecScalar = dyn_cast(op)) + } + if (auto vecScalar = dyn_cast(op)) { return verifyVecScalar(vecScalar, "pto.vmi.vmuls"); - if (auto vecScalar = dyn_cast(op)) + } + if (auto vecScalar = dyn_cast(op)) { return verifyVecScalar(vecScalar, "pto.vmi.vmaxs"); - if (auto vecScalar = dyn_cast(op)) + } + if (auto vecScalar = dyn_cast(op)) { return verifyVecScalar(vecScalar, "pto.vmi.vmins"); - if (auto vecScalar = dyn_cast(op)) + } + if (auto vecScalar = dyn_cast(op)) { return verifyVecScalar(vecScalar, "pto.vmi.vshls"); - if (auto vecScalar = dyn_cast(op)) + } + if (auto vecScalar = dyn_cast(op)) { return verifyVecScalar(vecScalar, "pto.vmi.vshrs"); + } if (auto vmull = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedVmullShape(vmull, &reason))) + if (succeeded(checkSupportedVmullShape(vmull, &reason))) { return WalkResult::advance(); + } vmull.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.vmull requires equal 64/128/256-lane i32/ui32 data " @@ -13585,8 +14301,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, op, "pto.vmi.ln", cast(ln.getResult().getType())); if (auto relu = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedReluShape(relu, &reason))) + if (succeeded(checkSupportedReluShape(relu, &reason))) { return WalkResult::advance(); + } relu.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.relu direct lowering requires physical vreg parts with " @@ -13621,8 +14338,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, cast(select.getResult().getType())); if (auto vselr = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedVselrShape(vselr, &reason))) + if (succeeded(checkSupportedVselrShape(vselr, &reason))) { return WalkResult::advance(); + } vselr.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.vselr supports only contiguous lane_stride=1 layouts " @@ -13635,8 +14353,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto cmpf = dyn_cast(op)) { WalkResult physical = emitMaskableUnsupported( op, "pto.vmi.cmpf", cast(cmpf.getLhs().getType())); - if (physical.wasInterrupted()) + if (physical.wasInterrupted()) { return physical; + } if (succeeded(checkSupportedComparePredicate( op, cmpf.getPredicate()))) return WalkResult::advance(); @@ -13646,8 +14365,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto cmpi = dyn_cast(op)) { WalkResult physical = emitMaskableUnsupported( op, "pto.vmi.cmpi", cast(cmpi.getLhs().getType())); - if (physical.wasInterrupted()) + if (physical.wasInterrupted()) { return physical; + } if (succeeded(checkSupportedComparePredicate( op, cmpi.getPredicate()))) return WalkResult::advance(); @@ -13669,8 +14389,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto compress = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedCompressShape(compress, &reason))) + if (succeeded(checkSupportedCompressShape(compress, &reason))) { return WalkResult::advance(); + } compress.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.compress lowers through pto.vsqz only for one " @@ -13681,8 +14402,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto compressStore = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedCompressStoreShape(compressStore, &reason))) + if (succeeded(checkSupportedCompressStoreShape(compressStore, &reason))) { return WalkResult::advance(); + } compressStore.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.compress_store lowers through pto.vsqz + pto.vstur " @@ -13783,8 +14505,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto reduce = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedGroupReduceShape(reduce, &reason))) + if (succeeded(checkSupportedGroupReduceShape(reduce, &reason))) { return WalkResult::advance(); + } reduce.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.group_reduce_minf lowers through pto.vcgmin/vmin for " @@ -13796,8 +14519,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto reduce = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedGroupReduceShape(reduce, &reason))) + if (succeeded(checkSupportedGroupReduceShape(reduce, &reason))) { return WalkResult::advance(); + } reduce.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.group_reduce_mini lowers through pto.vcgmin/vmin for " @@ -13865,8 +14589,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto fma = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedFmaShape(fma, &reason))) + if (succeeded(checkSupportedFmaShape(fma, &reason))) { return WalkResult::advance(); + } fma.emitError() << kVMIDiagUnsupportedPrefix << "pto.vmi.fma lowers through pto.vmula only for f16/bf16/f32 " @@ -13877,8 +14602,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto extf = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedExtFShape(extf, &reason))) + if (succeeded(checkSupportedExtFShape(extf, &reason))) { return WalkResult::advance(); + } extf.emitError() << kVMIDiagUnsupportedPrefix @@ -13892,8 +14618,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto truncf = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedTruncFShape(truncf, &reason))) + if (succeeded(checkSupportedTruncFShape(truncf, &reason))) { return WalkResult::advance(); + } truncf.emitError() << kVMIDiagUnsupportedPrefix @@ -13907,8 +14634,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto fptosi = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedFPToSIShape(fptosi, &reason))) + if (succeeded(checkSupportedFPToSIShape(fptosi, &reason))) { return WalkResult::advance(); + } fptosi.emitError() << kVMIDiagUnsupportedPrefix @@ -13920,8 +14648,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto fptoui = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedFPToUIShape(fptoui, &reason))) + if (succeeded(checkSupportedFPToUIShape(fptoui, &reason))) { return WalkResult::advance(); + } fptoui.emitError() << kVMIDiagUnsupportedPrefix @@ -13934,8 +14663,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto sitofp = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedSIToFPShape(sitofp, &reason))) + if (succeeded(checkSupportedSIToFPShape(sitofp, &reason))) { return WalkResult::advance(); + } sitofp.emitError() << kVMIDiagUnsupportedPrefix @@ -13947,8 +14677,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto extsi = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedExtSIShape(extsi, &reason))) + if (succeeded(checkSupportedExtSIShape(extsi, &reason))) { return WalkResult::advance(); + } extsi.emitError() << kVMIDiagUnsupportedPrefix @@ -13964,8 +14695,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto extui = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedExtUIShape(extui, &reason))) + if (succeeded(checkSupportedExtUIShape(extui, &reason))) { return WalkResult::advance(); + } extui.emitError() << kVMIDiagUnsupportedPrefix @@ -13981,8 +14713,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto trunci = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedTruncIShape(trunci, &reason))) + if (succeeded(checkSupportedTruncIShape(trunci, &reason))) { return WalkResult::advance(); + } trunci.emitError() << kVMIDiagUnsupportedPrefix @@ -13999,8 +14732,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto bitcast = dyn_cast(op)) { std::string reason; - if (succeeded(checkSupportedBitcastShape(bitcast, &reason))) + if (succeeded(checkSupportedBitcastShape(bitcast, &reason))) { return WalkResult::advance(); + } bitcast.emitError() << kVMIDiagUnsupportedPrefix @@ -14055,14 +14789,17 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto shuffle = dyn_cast(op)) { std::string reason; - if (succeeded(computeShuffleForwardingSourceParts(shuffle, &reason))) + if (succeeded(computeShuffleForwardingSourceParts(shuffle, &reason))) { return WalkResult::advance(); + } std::string splatReason; - if (succeeded(computeShuffleLane0SplatSourcePart(shuffle, &splatReason))) + if (succeeded(computeShuffleLane0SplatSourcePart(shuffle, &splatReason))) { return WalkResult::advance(); + } std::string vselrReason; - if (succeeded(computeShuffleVselrPlans(shuffle, &vselrReason))) + if (succeeded(computeShuffleVselrPlans(shuffle, &vselrReason))) { return WalkResult::advance(); + } shuffle.emitError() << kVMIDiagUnsupportedPrefix @@ -14075,8 +14812,9 @@ verifySupportedVMIToVPTOOps(ModuleOp module, if (auto constantMask = dyn_cast(op)) { std::string reason; - if (succeeded(computeConstantMaskMaterialization(constantMask, &reason))) + if (succeeded(computeConstantMaskMaterialization(constantMask, &reason))) { return WalkResult::advance(); + } constantMask.emitError() << kVMIDiagUnsupportedPrefix diff --git a/lib/PTO/Transforms/VPTOBufferMaterialization.cpp b/lib/PTO/Transforms/VPTOBufferMaterialization.cpp index 33abcc91e4..de99742385 100644 --- a/lib/PTO/Transforms/VPTOBufferMaterialization.cpp +++ b/lib/PTO/Transforms/VPTOBufferMaterialization.cpp @@ -18,11 +18,13 @@ namespace { static AddressSpaceAttr getNormalizedPtrMemorySpace(Attribute memorySpace, MLIRContext *context) { - if (auto addrSpace = dyn_cast_or_null(memorySpace)) + if (auto addrSpace = dyn_cast_or_null(memorySpace)) { return addrSpace; - if (auto intAttr = dyn_cast_or_null(memorySpace)) + } + if (auto intAttr = dyn_cast_or_null(memorySpace)) { return AddressSpaceAttr::get(context, static_cast(intAttr.getInt())); + } return AddressSpaceAttr::get(context, AddressSpace::GM); } @@ -31,8 +33,9 @@ static Value materializeMemRefView(Value value, ArrayRef shape, PatternRewriter &rewriter, Location loc) { auto memrefType = MemRefType::get(shape, elementType, AffineMap(), memorySpace); - if (value.getType() == memrefType) + if (value.getType() == memrefType) { return value; + } return rewriter .create( loc, TypeRange(ArrayRef{memrefType}), value) @@ -41,12 +44,14 @@ static Value materializeMemRefView(Value value, ArrayRef shape, static Value materializeTileBufferView(Value value, PatternRewriter &rewriter, Location loc) { - if (isa(value.getType())) + if (isa(value.getType())) { return value; + } auto tileType = dyn_cast(value.getType()); - if (!tileType) + if (!tileType) { return {}; + } return materializeMemRefView(value, tileType.getShape(), tileType.getElementType(), @@ -58,21 +63,23 @@ static Value materializeTileBufferView(Value value, PatternRewriter &rewriter, Value materializeBufferPointer(Value value, Type elementType, Attribute memorySpace, PatternRewriter &rewriter, Location loc) { - if (!value) + if (!value) { return {}; + } auto ptrMemorySpace = getNormalizedPtrMemorySpace(memorySpace, rewriter.getContext()); auto ptrType = PtrType::get(rewriter.getContext(), elementType, ptrMemorySpace); - if (value.getType() == ptrType) + if (value.getType() == ptrType) { return value; + } Value memrefValue = materializeTileBufferView(value, rewriter, loc); auto memrefType = dyn_cast_or_null(memrefValue.getType()); - if (!memrefValue || !memrefType) + if (!memrefValue || !memrefType) { return {}; + } return rewriter.create(loc, ptrType, memrefValue).getResult(); } - } // namespace mlir::pto diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index db5f9a344e..53adc3d39c 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -68,16 +68,21 @@ static Type getElementTypeFromVectorLike(Type type); static std::optional getElementCountFromVectorLike(Type type); static Type getLowPrecisionLLVMType(Type type, MLIRContext *context) { - if (pto::isPTOHiFloat8Type(type)) + if (pto::isPTOHiFloat8Type(type)) { return LLVM::LLVMHiFloat8Type::get(context); - if (isa(type)) + } + if (isa(type)) { return LLVM::LLVMFloat4E1M2x2Type::get(context); - if (isa(type)) + } + if (isa(type)) { return LLVM::LLVMFloat4E2M1x2Type::get(context); - if (pto::isPTOFloat8E4M3LikeType(type)) + } + if (pto::isPTOFloat8E4M3LikeType(type)) { return LLVM::LLVMFloat8E4M3Type::get(context); - if (pto::isPTOFloat8E5M2LikeType(type)) + } + if (pto::isPTOFloat8E5M2LikeType(type)) { return LLVM::LLVMFloat8E5M2Type::get(context); + } return {}; } @@ -91,20 +96,23 @@ static Type normalizePayloadTypeForLLVMLowering(Type type, Builder &builder) { if (pto::isPTOHiFloat8x2Type(type)) return getLLVMCompatibleVectorType( {2}, LLVM::LLVMHiFloat8Type::get(builder.getContext())); - if (Type lowpType = getLowPrecisionLLVMType(type, builder.getContext())) + if (Type lowpType = getLowPrecisionLLVMType(type, builder.getContext())) { return lowpType; + } if (auto intType = dyn_cast(type)) { - if (!intType.isSignless()) + if (!intType.isSignless()) { return builder.getIntegerType(intType.getWidth()); + } return type; } if (auto vecType = dyn_cast(type)) { Type normalizedElement = normalizePayloadTypeForLLVMLowering(vecType.getElementType(), builder); - if (normalizedElement == vecType.getElementType()) + if (normalizedElement == vecType.getElementType()) { return type; + } return getLLVMCompatibleVectorType(vecType.getShape(), normalizedElement, vecType.getScalableDims()); } @@ -114,10 +122,12 @@ static Type normalizePayloadTypeForLLVMLowering(Type type, Builder &builder) { static Type normalizeGEPElementTypeForLLVMLowering(Type type, Builder &builder) { - if (pto::isPTOHiFloat8x2Type(type)) + if (pto::isPTOHiFloat8x2Type(type)) { return builder.getI16Type(); - if (pto::isPTOLowPrecisionType(type)) + } + if (pto::isPTOLowPrecisionType(type)) { return builder.getI8Type(); + } if (isa(type)) @@ -127,8 +137,9 @@ static Type normalizeGEPElementTypeForLLVMLowering(Type type, Type normalizedElement = normalizeGEPElementTypeForLLVMLowering(vecType.getElementType(), builder); - if (normalizedElement == vecType.getElementType()) + if (normalizedElement == vecType.getElementType()) { return normalizePayloadTypeForLLVMLowering(type, builder); + } return getLLVMCompatibleVectorType(vecType.getShape(), normalizedElement, vecType.getScalableDims()); } @@ -143,12 +154,15 @@ static Type convertVPTOType(Type type, Builder &builder) { return getLLVMCompatibleVectorType({vecType.getElementCount()}, elementType); } - if (isa(type)) + if (isa(type)) { return VectorType::get({256}, builder.getI1Type()); - if (isa(type)) + } + if (isa(type)) { return VectorType::get({32}, builder.getI8Type()); - if (isa(type)) + } + if (isa(type)) { return LLVM::LLVMPointerType::get(builder.getContext()); + } if (auto ptrType = dyn_cast(type)) { return LLVM::LLVMPointerType::get( builder.getContext(), @@ -160,37 +174,47 @@ static Type convertVPTOType(Type type, Builder &builder) { static unsigned getNaturalByteAlignment(Type type) { if (auto vecType = dyn_cast(type)) { unsigned elemAlign = getNaturalByteAlignment(vecType.getElementType()); - if (!elemAlign) + if (!elemAlign) { return 0; + } int64_t elems = 1; - for (int64_t dim : vecType.getShape()) + for (int64_t dim : vecType.getShape()) { elems *= dim; + } return elemAlign * static_cast(elems); } - if (auto intType = dyn_cast(type)) - return llvm::divideCeil(unsigned(intType.getWidth()), 8u); - if (pto::isPTOHiFloat8x2Type(type)) + if (auto intType = dyn_cast(type)) { + return llvm::divideCeil(static_cast(intType.getWidth()), 8U); + } + if (pto::isPTOHiFloat8x2Type(type)) { return 2; - if (pto::isPTOLowPrecisionType(type)) + } + if (pto::isPTOLowPrecisionType(type)) { return 1; - if (type.isF16() || type.isBF16()) + } + if (type.isF16() || type.isBF16()) { return 2; - if (type.isF32()) + } + if (type.isF32()) { return 4; - if (type.isF64()) + } + if (type.isF64()) { return 8; + } return 0; } static bool hasVPTOConvertibleType(Type type) { - if (!type) + if (!type) { return false; + } if (isa(type) || pto::isPTOLowPrecisionType(type)) return true; - if (auto vecType = dyn_cast(type)) + if (auto vecType = dyn_cast(type)) { return hasVPTOConvertibleType(vecType.getElementType()); + } return false; } @@ -200,8 +224,9 @@ static bool hasVPTOConvertibleType(TypeRange types) { static Value materializeVPTOCast(OpBuilder &builder, Type resultType, ValueRange inputs, Location loc) { - if (inputs.size() != 1) + if (inputs.size() != 1) { return {}; + } return builder .create(loc, TypeRange{resultType}, inputs) .getResult(0); @@ -242,8 +267,9 @@ static LLVM::LLVMStructType getVPTOStructStorageType(pto::StructType structType, if (!frame.materialize) { worklist.push_back({frame.type, true}); for (Type fieldType : frame.type.getFieldTypes()) { - if (auto nestedStruct = dyn_cast(fieldType)) + if (auto nestedStruct = dyn_cast(fieldType)) { worklist.push_back({nestedStruct, false}); + } } continue; } @@ -271,18 +297,21 @@ getVPTOStructFieldAddress(ConversionPatternRewriter &rewriter, Location loc, Value address = root; pto::StructType currentType = rootType; for (auto [depth, index] : llvm::enumerate(path)) { - if (index < 0 || index >= static_cast(currentType.getNumFields())) + if (index < 0 || index >= static_cast(currentType.getNumFields())) { return failure(); + } Type storageType = getVPTOStructStorageType(currentType, rewriter); address = rewriter.create( loc, pointerType, storageType, address, ArrayRef{0, static_cast(index)}); Type fieldType = currentType.getFieldType(static_cast(index)); - if (depth + 1 == path.size()) + if (depth + 1 == path.size()) { continue; + } auto nestedStruct = dyn_cast(fieldType); - if (!nestedStruct) + if (!nestedStruct) { return failure(); + } currentType = nestedStruct; } return address; @@ -367,10 +396,12 @@ static Value getI32Constant(OpBuilder &builder, Location loc, uint64_t value) { } static bool isMxElementType(Type ty) { - if (auto floatType = dyn_cast(ty)) + if (auto floatType = dyn_cast(ty)) { return floatType.getWidth() == 8; - if (isa(ty)) + } + if (isa(ty)) { return true; + } std::string typeText; llvm::raw_string_ostream os(typeText); ty.print(os); @@ -379,10 +410,12 @@ static bool isMxElementType(Type ty) { } static std::string getMadMxElementFragment(Type type) { - if (type.isF16()) + if (type.isF16()) { return "f16"; - if (type.isBF16()) + } + if (type.isBF16()) { return "bf16"; + } std::string typeText; llvm::raw_string_ostream os(typeText); @@ -390,16 +423,21 @@ static std::string getMadMxElementFragment(Type type) { os.flush(); std::string lower = StringRef(typeText).lower(); - if (StringRef(lower).contains("e4m3")) + if (StringRef(lower).contains("e4m3")) { return "e4m3"; - if (StringRef(lower).contains("e5m2")) + } + if (StringRef(lower).contains("e5m2")) { return "e5m2"; - if (StringRef(lower).contains("hif4")) + } + if (StringRef(lower).contains("hif4")) { return "hif4"; - if (StringRef(lower).contains("e2m1x2")) + } + if (StringRef(lower).contains("e2m1x2")) { return "e2m1x2"; - if (StringRef(lower).contains("e1m2x2")) + } + if (StringRef(lower).contains("e1m2x2")) { return "e1m2x2"; + } return {}; } @@ -407,8 +445,9 @@ static FailureOr buildMadMxCalleeName(MLIRContext *context, Type lhsElem, Type rhsElem) { std::string lhs = getMadMxElementFragment(lhsElem); std::string rhs = getMadMxElementFragment(rhsElem); - if (lhs.empty() || rhs.empty()) + if (lhs.empty() || rhs.empty()) { return failure(); + } return StringAttr::get(context, "llvm.hivm.MMAD.MX." + lhs + rhs).getValue(); } @@ -418,19 +457,25 @@ static bool isSignedOrSignlessInteger(IntegerType intType, unsigned width) { } static std::string getMadRhsFragment(Type type) { - if (type.isF16()) + if (type.isF16()) { return "f16"; - if (type.isBF16()) + } + if (type.isBF16()) { return "bf16"; - if (type.isF32()) + } + if (type.isF32()) { return "f32"; + } if (auto intType = dyn_cast(type)) { - if (isSignedOrSignlessInteger(intType, 4)) + if (isSignedOrSignlessInteger(intType, 4)) { return "s4"; - if (isSignedOrSignlessInteger(intType, 8)) + } + if (isSignedOrSignlessInteger(intType, 8)) { return "s8"; - if (intType.isUnsigned() && intType.getWidth() == 2) + } + if (intType.isUnsigned() && intType.getWidth() == 2) { return "u2"; + } } std::string typeText; @@ -438,8 +483,9 @@ static std::string getMadRhsFragment(Type type) { type.print(os); os.flush(); std::string lower = StringRef(typeText).lower(); - if (StringRef(lower).contains("e8m0")) + if (StringRef(lower).contains("e8m0")) { return "e8m0"; + } return {}; } @@ -452,13 +498,16 @@ static bool isMadE5M2ElementType(Type type) { } static std::string getMadDstFragment(Type type) { - if (type.isF16()) + if (type.isF16()) { return "f16"; - if (type.isF32()) + } + if (type.isF32()) { return "f32"; + } if (auto intType = dyn_cast(type)) { - if (isSignedOrSignlessInteger(intType, 32)) + if (isSignedOrSignlessInteger(intType, 32)) { return "s32"; + } } return {}; } @@ -6942,7 +6991,7 @@ class LowerVldsx2OpPattern final : public OpConversionPattern { SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), resultTypes)) || - resultTypes.size() != (usePostIntrinsic ? 3u : 2u)) { + resultTypes.size() != (usePostIntrinsic ? 3U : 2U)) { return rewriter.notifyMatchFailure(op, "failed to convert vldsx2 result types"); } @@ -7008,7 +7057,7 @@ class LowerVsldbOpPattern final : public OpConversionPattern { SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), resultTypes)) || - resultTypes.size() != (usePostIntrinsic ? 2u : 1u)) + resultTypes.size() != (usePostIntrinsic ? 2U : 1U)) return rewriter.notifyMatchFailure(op, "failed to convert vsldb result type"); Type callResultType = getPayloadABIType( @@ -7121,7 +7170,7 @@ class LowerVldusOpPattern final : public OpConversionPattern { bool usePostIntrinsic = static_cast(op.getUpdatedBase()); if (!sourceType || failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), resultTypes)) || - resultTypes.size() != (usePostIntrinsic ? 3u : 2u) || + resultTypes.size() != (usePostIntrinsic ? 3U : 2U) || adaptor.getAlign().getType() != resultTypes[1] || (usePostIntrinsic && resultTypes[2] != adaptor.getSource().getType())) { return rewriter.notifyMatchFailure(op, @@ -7223,7 +7272,7 @@ class LowerSprStoreOpPattern final : public OpConversionPattern { SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), resultTypes)) || - resultTypes.size() != (usePostIntrinsic ? 1u : 0u)) + resultTypes.size() != (usePostIntrinsic ? 1U : 0U)) return rewriter.notifyMatchFailure( op, "failed to convert spr store result types"); @@ -7512,7 +7561,7 @@ class LowerVstusOpPattern final : public OpConversionPattern { "failed to convert vstus result types"); bool usePostIntrinsic = static_cast(op.getBaseOut()); auto baseType = dyn_cast(adaptor.getBase().getType()); - if (!baseType || resultTypes.size() != (usePostIntrinsic ? 2u : 1u) || + if (!baseType || resultTypes.size() != (usePostIntrinsic ? 2U : 1U) || adaptor.getAlignIn().getType() != resultTypes[0] || (usePostIntrinsic && resultTypes[1] != adaptor.getBase().getType())) { return rewriter.notifyMatchFailure(op, @@ -7647,7 +7696,7 @@ class LowerVstasOpPattern final : public OpConversionPattern { SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), resultTypes)) || - resultTypes.size() != (usePostIntrinsic ? 1u : 0u)) + resultTypes.size() != (usePostIntrinsic ? 1U : 0U)) return rewriter.notifyMatchFailure( op, "failed to convert vstas result types"); @@ -8323,7 +8372,7 @@ class LowerPredicateStoreOpPattern final : public OpConversionPattern { SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), resultTypes)) || - resultTypes.size() != (usePostIntrinsic ? 1u : 0u)) + resultTypes.size() != (usePostIntrinsic ? 1U : 0U)) return rewriter.notifyMatchFailure( op, "failed to convert predicate-store result types"); @@ -8373,7 +8422,7 @@ class LowerPredicateLoadOpPattern final : public OpConversionPattern { SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), resultTypes)) || - resultTypes.size() != (usePostIntrinsic ? 2u : 1u)) + resultTypes.size() != (usePostIntrinsic ? 2U : 1U)) return rewriter.notifyMatchFailure( op, "failed to convert predicate-load result types"); if (!llvmSourceType) @@ -8962,20 +9011,20 @@ class LowerPipeEventDynSyncOpPattern final : public OpConversionPattern StringRef calleeName = buildSyncCallee(op.getContext()); Value srcValue = getI64Constant(rewriter, op.getLoc(), *src); Value dstValue = getI64Constant(rewriter, op.getLoc(), *dst); - + Value eventIdValue = adaptor.getEventId(); if (!eventIdValue) return rewriter.notifyMatchFailure(op, "missing event_id operand"); - + Value eventValue = eventIdValue; - + while (eventValue.getDefiningOp()) { auto unrealizedCast = dyn_cast(eventValue.getDefiningOp()); if (!unrealizedCast || unrealizedCast.getInputs().size() != 1) break; eventValue = unrealizedCast.getInputs()[0]; } - + if (eventValue.getType().isIndex()) { eventValue = rewriter.create(op.getLoc(), rewriter.getI64Type(), @@ -8989,7 +9038,7 @@ class LowerPipeEventDynSyncOpPattern final : public OpConversionPattern } else { return rewriter.notifyMatchFailure(op, "unexpected event_id type"); } - + auto funcType = rewriter.getFunctionType( TypeRange{rewriter.getI64Type(), rewriter.getI64Type(), rewriter.getI64Type()}, diff --git a/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp b/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp index 2bb75efcb4..cf8f16032a 100644 --- a/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp +++ b/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp @@ -38,18 +38,21 @@ enum class DmaArch { A2A3, A5 }; constexpr uint64_t kMxScaleAddressShift = 4; static DmaArch getDmaArch(ModuleOp mod) { - if (!mod) + if (!mod) { return DmaArch::A2A3; + } auto arch = mod->getAttrOfType("pto.target_arch"); - if (arch && arch.getValue() == "a5") + if (arch && arch.getValue() == "a5") { return DmaArch::A5; + } return DmaArch::A2A3; } static pto::AddressSpaceAttr getPointerMemorySpace(Attribute memorySpace, MLIRContext *ctx) { - if (auto addrSpace = dyn_cast_or_null(memorySpace)) + if (auto addrSpace = dyn_cast_or_null(memorySpace)) { return addrSpace; + } if (auto intAttr = dyn_cast_or_null(memorySpace)) return pto::AddressSpaceAttr::get( ctx, static_cast(intAttr.getInt())); @@ -58,48 +61,57 @@ static pto::AddressSpaceAttr getPointerMemorySpace(Attribute memorySpace, static bool hasZeroReinterpretOffset(memref::ReinterpretCastOp op) { for (int64_t offset : op.getStaticOffsets()) { - if (ShapedType::isDynamic(offset) || offset != 0) + if (ShapedType::isDynamic(offset) || offset != 0) { return false; + } } return true; } static Value materializeBufferPointer(Value value, PatternRewriter &rewriter, Location loc) { - if (!value) + if (!value) { return {}; + } if (auto ptrType = dyn_cast(value.getType())) { if (auto cast = value.getDefiningOp()) { - if (cast->getNumOperands() != 1 || cast->getNumResults() != 1) + if (cast->getNumOperands() != 1 || cast->getNumResults() != 1) { return {}; + } Value basePtr = materializeBufferPointer(cast.getOperand(0), rewriter, loc); - if (!basePtr) + if (!basePtr) { return {}; - if (basePtr.getType() == ptrType) + } + if (basePtr.getType() == ptrType) { return basePtr; + } return rewriter.create(loc, ptrType, basePtr).getResult(); } return value; } if (auto cast = value.getDefiningOp()) { - if (cast->getNumOperands() != 1 || cast->getNumResults() != 1) + if (cast->getNumOperands() != 1 || cast->getNumResults() != 1) { return {}; + } return materializeBufferPointer(cast.getOperand(0), rewriter, loc); } - if (auto cast = value.getDefiningOp()) + if (auto cast = value.getDefiningOp()) { return materializeBufferPointer(cast.getSource(), rewriter, loc); + } - if (auto cast = value.getDefiningOp()) + if (auto cast = value.getDefiningOp()) { return materializeBufferPointer(cast.getSource(), rewriter, loc); + } if (auto cast = value.getDefiningOp()) { auto resultType = dyn_cast(value.getType()); - if (!resultType) + if (!resultType) { return {}; + } if (hasZeroReinterpretOffset(cast)) { Value basePtr = materializeBufferPointer(cast.getSource(), rewriter, loc); @@ -108,16 +120,18 @@ static Value materializeBufferPointer(Value value, PatternRewriter &rewriter, rewriter.getContext(), resultType.getElementType(), getPointerMemorySpace(resultType.getMemorySpace(), rewriter.getContext())); - if (basePtr.getType() == ptrType) + if (basePtr.getType() == ptrType) { return basePtr; + } return rewriter.create(loc, ptrType, basePtr).getResult(); } } } auto memrefType = dyn_cast(value.getType()); - if (!memrefType) + if (!memrefType) { return {}; + } auto ptrType = pto::PtrType::get(rewriter.getContext(), memrefType.getElementType(), @@ -127,18 +141,21 @@ static Value materializeBufferPointer(Value value, PatternRewriter &rewriter, } static Type getBufferElementType(Type type) { - if (auto ptrType = dyn_cast(type)) + if (auto ptrType = dyn_cast(type)) { return ptrType.getElementType(); - if (auto memrefType = dyn_cast(type)) + } + if (auto memrefType = dyn_cast(type)) { return memrefType.getElementType(); + } return {}; } static Value offsetBufferPointer(Value basePtr, Type elementType, Value elementOffset, PatternRewriter &rewriter, Location loc) { - if (!basePtr) + if (!basePtr) { return {}; + } Value offsetIndex = elementOffset; if (!offsetIndex.getType().isIndex()) @@ -156,8 +173,9 @@ static bool isKnownOne(Value value) { } static bool shouldRestoreDmaLoopSize(Value loop1Count, Value loop2Count) { - if (!loop1Count) + if (!loop1Count) { return false; + } return !isKnownOne(loop1Count) || !isKnownOne(loop2Count); } @@ -174,17 +192,21 @@ static SmallVector collectLoopConfigs(ValueRange counts, static Value offsetPointerByBytes(Value basePtr, Value byteOffset, PatternRewriter &rewriter, Location loc) { - if (!basePtr) + if (!basePtr) { return {}; + } Value basePtrValue = materializeBufferPointer(basePtr, rewriter, loc); auto ptrType = dyn_cast_or_null(basePtrValue.getType()); - if (!ptrType) + if (!ptrType) { return {}; + } APInt constOffset; - if (matchPattern(byteOffset, m_ConstantInt(&constOffset)) && constOffset.isZero()) + if (matchPattern(byteOffset, m_ConstantInt(&constOffset)) && + constOffset.isZero()) { return basePtrValue; + } auto bytePtrType = pto::PtrType::get(rewriter.getContext(), rewriter.getI8Type(), @@ -204,35 +226,44 @@ static Value offsetPointerByBytes(Value basePtr, Value byteOffset, [[maybe_unused]] static Value materializeFpcValue(Value fpc, PatternRewriter &rewriter, Location loc) { - if (!fpc) + if (!fpc) { return {}; - if (fpc.getType().isInteger(64)) + } + if (fpc.getType().isInteger(64)) { return fpc; - if (isa(fpc.getType())) + } + if (isa(fpc.getType())) { return rewriter.create(loc, rewriter.getI64Type(), fpc); + } return {}; } static Value materializeI64Value(Value value, PatternRewriter &rewriter, Location loc) { - if (!value) + if (!value) { return {}; - if (value.getType().isInteger(64)) + } + if (value.getType().isInteger(64)) { return value; - if (auto intType = dyn_cast(value.getType())) + } + if (auto intType = dyn_cast(value.getType())) { return rewriter.create(loc, rewriter.getI64Type(), value); - if (isa(value.getType())) + } + if (isa(value.getType())) { return rewriter.create(loc, rewriter.getI64Type(), value); + } return {}; } static Value materializeAccStoreScalarPayload(Value value, PatternRewriter &rewriter, Location loc) { - if (!value) + if (!value) { return {}; - if (Value raw = materializeI64Value(value, rewriter, loc)) + } + if (Value raw = materializeI64Value(value, rewriter, loc)) { return raw; + } Type type = value.getType(); Value f32Value = value; @@ -249,8 +280,9 @@ static Value materializeAccStoreScalarPayload(Value value, static Value materializeAccStoreClipPayload(Value value, Type destinationElementType, PatternRewriter &rewriter, Location loc) { - if (!value) + if (!value) { return {}; + } if (value.getType().isF16()) { Value bitsI16 = @@ -259,8 +291,9 @@ static Value materializeAccStoreClipPayload(Value value, Type destinationElement } auto intType = dyn_cast(value.getType()); - if (!intType) + if (!intType) { return {}; + } Value widened; if (auto dstIntType = dyn_cast(destinationElementType); @@ -283,8 +316,9 @@ static Value deriveMxScaleDestination(Value dataDestination, PatternRewriter &rewriter, Location loc) { auto ptrType = dyn_cast(dataDestination.getType()); - if (!ptrType) + if (!ptrType) { return {}; + } Value dataAddress = rewriter.create( loc, rewriter.getI64Type(), dataDestination); @@ -334,8 +368,9 @@ static Value buildAccStoreFpcValue(Location loc, Value preQuant, case pto::AccStoreQuantPreMode::QF322F16PreVec: case pto::AccStoreQuantPreMode::QF322BF16PreVec: case pto::AccStoreQuantPreMode::QS322BF16PreVec: - if (Value quantPtr = materializeI64Value(preQuant, rewriter, loc)) + if (Value quantPtr = materializeI64Value(preQuant, rewriter, loc)) { quantAddr = encodeFixpipeBufferAddr(quantPtr, /*unitShift=*/7); + } break; default: break; @@ -344,12 +379,14 @@ static Value buildAccStoreFpcValue(Location loc, Value preQuant, Value reluAddr; if (preReluMode && *preReluMode == pto::ReluPreMode::VectorRelu) { - if (Value reluPtr = materializeI64Value(preRelu, rewriter, loc)) + if (Value reluPtr = materializeI64Value(preRelu, rewriter, loc)) { reluAddr = encodeFixpipeBufferAddr(reluPtr, /*unitShift=*/6); + } } - if (!quantAddr && !reluAddr) + if (!quantAddr && !reluAddr) { return {}; + } Value mask = getI64Constant(loc, rewriter, 0xff); Value fpc = getI64Constant(loc, rewriter, 0); @@ -400,12 +437,16 @@ static void configureAccStoreScalarPreOps(Location loc, Value preQuant, if (preQuantMode && *preQuantMode != pto::AccStoreQuantPreMode::NoConvert && !isVectorQuantMode(*preQuantMode)) { - if (Value quantValue = materializeAccStoreScalarPayload(preQuant, rewriter, loc)) + if (Value quantValue = + materializeAccStoreScalarPayload(preQuant, rewriter, loc)) { rewriter.create(loc, quantValue); + } } if (preReluMode && *preReluMode == pto::ReluPreMode::ScalarRelu) { - if (Value reluAlpha = materializeAccStoreScalarPayload(preRelu, rewriter, loc)) + if (Value reluAlpha = + materializeAccStoreScalarPayload(preRelu, rewriter, loc)) { rewriter.create(loc, reluAlpha); + } } if (clipValue) { if (Value clip = materializeAccStoreClipPayload(clipValue, @@ -420,8 +461,9 @@ static Value configureAccStoreCtrl(Location loc, bool allowAtomic, std::optional atomicOp, std::optional satMode, PatternRewriter &rewriter) { - if ((!allowAtomic || !atomicType || !atomicOp) && !satMode) + if ((!allowAtomic || !atomicType || !atomicOp) && !satMode) { return {}; + } Value originalCtrl = rewriter.create(loc); Value ctrl = originalCtrl; @@ -480,22 +522,27 @@ static Value packLoopSize(Location loc, Value loop2, Value loop1, static Value castIntegerLikeTo(Location loc, Value value, Type targetType, PatternRewriter &rewriter) { - if (value.getType() == targetType) + if (value.getType() == targetType) { return value; + } auto targetInt = dyn_cast(targetType); - if (value.getType().isIndex() && targetInt) + if (value.getType().isIndex() && targetInt) { return rewriter.create(loc, targetType, value); + } if (auto sourceInt = dyn_cast(value.getType())) { if (targetInt) { - if (sourceInt.getWidth() < targetInt.getWidth()) + if (sourceInt.getWidth() < targetInt.getWidth()) { return rewriter.create(loc, targetType, value); - if (sourceInt.getWidth() > targetInt.getWidth()) + } + if (sourceInt.getWidth() > targetInt.getWidth()) { return rewriter.create(loc, targetType, value); + } return value; } - if (targetType.isIndex()) + if (targetType.isIndex()) { return rewriter.create(loc, targetType, value); + } } return {}; @@ -510,8 +557,9 @@ static FailureOr packMadXt(Location loc, Value m, Value n, Value k, Value mI64 = castIntegerLikeTo(loc, m, i64Ty, rewriter); Value nI64 = castIntegerLikeTo(loc, n, i64Ty, rewriter); Value kI64 = castIntegerLikeTo(loc, k, i64Ty, rewriter); - if (!mI64 || !nI64 || !kI64) + if (!mI64 || !nI64 || !kI64) { return failure(); + } auto constant = [&](uint64_t value) -> Value { return rewriter.create(loc, value, 64); @@ -531,20 +579,24 @@ static FailureOr packMadXt(Location loc, Value m, Value n, Value k, *unitFlagMode == pto::MadUnitFlagMode::CheckOnly ? 2 : 3; xt = bitOr(xt, shl(constant(unitFlagCtrl), 55)); } - if (disableGemv) + if (disableGemv) { xt = bitOr(xt, shl(constant(1), 61)); - if (cmatrixSource) + } + if (cmatrixSource) { xt = bitOr(xt, shl(constant(1), 62)); - if (cmatrixInit) + } + if (cmatrixInit) { xt = bitOr(xt, shl(constant(1), 63)); + } return xt; } static Value setCtrlBit(Location loc, Value ctrl, unsigned bitIndex, bool value, PatternRewriter &rewriter) { Value bit = rewriter.create(loc, bitIndex, 64); - if (value) + if (value) { return rewriter.create(loc, ctrl, bit).getResult(); + } return rewriter.create(loc, ctrl, bit).getResult(); } @@ -563,9 +615,10 @@ static Value buildMadSemanticCtrl(Location loc, Value ctrl, ctrl = setCtrlBit(loc, ctrl, 46, false, rewriter); ctrl = setCtrlBit(loc, ctrl, 47, false, rewriter); } - if (satMode) + if (satMode) { ctrl = setCtrlBit(loc, ctrl, 48, *satMode == pto::MadSatMode::NoSat, rewriter); + } ctrl = setCtrlBit(loc, ctrl, 51, hasNDir, rewriter); return ctrl; } @@ -803,8 +856,9 @@ deriveLoadCbufToCbControl(Location loc, Value k, Value n, Type elementType, Value mStart, Value kStart, bool transpose, PatternRewriter &rewriter) { unsigned elemBitWidth = pto::getPTOStorageElemBitWidth(elementType); - if (elemBitWidth == 0 || (elemBitWidth % 8) != 0) + if (elemBitWidth == 0 || (elemBitWidth % 8) != 0) { return failure(); + } uint64_t elemBytes = elemBitWidth / 8; bool isFp4Packed = pto::isPTOFloat4PackedType(elementType); @@ -856,8 +910,9 @@ deriveLoadCbufToCaControl(Location loc, Value m, Value k, Type elementType, Value mStart, Value kStart, bool transpose, PatternRewriter &rewriter) { unsigned elemBitWidth = pto::getPTOStorageElemBitWidth(elementType); - if (elemBitWidth == 0 || (elemBitWidth % 8) != 0) + if (elemBitWidth == 0 || (elemBitWidth % 8) != 0) { return failure(); + } uint64_t elemBytes = elemBitWidth / 8; bool isFp4Packed = pto::isPTOFloat4PackedType(elementType); @@ -909,8 +964,9 @@ deriveLoadCbufToCaMxControl(Location loc, Value m, Value k, Type elementType, Value startRow, Value startCol, PatternRewriter &rewriter) { unsigned elemBitWidth = pto::getPTOStorageElemBitWidth(elementType); - if (elemBitWidth == 0 || (elemBitWidth % 8) != 0) + if (elemBitWidth == 0 || (elemBitWidth % 8) != 0) { return failure(); + } uint64_t elemBytes = elemBitWidth / 8; auto constant = [&](uint64_t value) -> Value { @@ -938,8 +994,9 @@ deriveLoadCbufToCbMxControl(Location loc, Value k, Value n, Type elementType, Value startRow, Value startCol, PatternRewriter &rewriter) { unsigned elemBitWidth = pto::getPTOStorageElemBitWidth(elementType); - if (elemBitWidth == 0 || (elemBitWidth % 8) != 0) + if (elemBitWidth == 0 || (elemBitWidth % 8) != 0) { return failure(); + } uint64_t elemBytes = elemBitWidth / 8; auto constant = [&](uint64_t value) -> Value { @@ -1011,8 +1068,9 @@ struct ExpandUvldPattern : public OpRewritePattern { LogicalResult matchAndRewrite(pto::UvldOp op, PatternRewriter &rewriter) const override { auto vecType = dyn_cast(op.getResult().getType()); - if (!vecType) + if (!vecType) { return failure(); + } Value basePtr = materializeBufferPointer(op.getSource(), rewriter, op.getLoc()); if (!basePtr) @@ -1035,8 +1093,9 @@ struct ExpandUvldPattern : public OpRewritePattern { enum class MadRawKind { Ordinary, OrdinaryBias, Mx, MxBias }; static MadRawKind deriveMadRawKind(pto::MadSemanticOpInterface op) { - if (op.isMadMxFamily()) + if (op.isMadMxFamily()) { return op.hasBiasOperand() ? MadRawKind::MxBias : MadRawKind::Mx; + } return op.hasBiasOperand() ? MadRawKind::OrdinaryBias : MadRawKind::Ordinary; } @@ -1087,8 +1146,9 @@ static LogicalResult lowerMadSemanticOp(pto::MadSemanticOpInterface op, satMode = satModeAttr.getValue(); bool isHif8 = false; - if (auto lhsPtr = dyn_cast(op.getLhs().getType())) + if (auto lhsPtr = dyn_cast(op.getLhs().getType())) { isHif8 = pto::isPTOHiFloat8Type(lhsPtr.getElementType()); + } Location loc = op->getLoc(); Value ctrlSaved = rewriter.create(loc).getResult(); @@ -1100,11 +1160,13 @@ static LogicalResult lowerMadSemanticOp(pto::MadSemanticOpInterface op, packMadXt(loc, op.getM(), op.getN(), op.getK(), unitFlagMode, op.getDisableGemv(), op.initializesAccumulatorWithBias(), op.initializesAccumulatorWithZero(), rewriter); - if (failed(xt)) + if (failed(xt)) { return rewriter.notifyMatchFailure(op, "failed to pack mad xt"); + } - if (failed(emitMadRawOp(op, deriveMadRawKind(op), *xt, rewriter))) + if (failed(emitMadRawOp(op, deriveMadRawKind(op), *xt, rewriter))) { return rewriter.notifyMatchFailure(op, "failed to emit mad raw op"); + } rewriter.create(loc, ctrlSaved); rewriter.eraseOp(op); @@ -1120,8 +1182,9 @@ class ExpandMadSemanticPattern final : public OpRewritePattern { LogicalResult matchAndRewrite(SemanticOp op, PatternRewriter &rewriter) const override { auto semantic = dyn_cast(op.getOperation()); - if (!semantic) + if (!semantic) { return failure(); + } return lowerMadSemanticOp(semantic, rewriter); } }; @@ -1174,18 +1237,21 @@ struct ExpandDmaLoadPattern : public OpRewritePattern { } Value leftPadding = op.getLeftPaddingCount(); - if (!leftPadding) + if (!leftPadding) { leftPadding = rewriter.create(loc, 0, 64); + } Value rightPadding = op.getRightPaddingCount(); - if (!rightPadding) + if (!rightPadding) { rightPadding = rewriter.create(loc, 0, 64); + } Value dataSelect = rewriter.create( loc, rewriter.getI1Type(), rewriter.getBoolAttr(static_cast(op.getPadValue()))); bool hasPad = static_cast(op.getPadValue()); - if (Value padValue = op.getPadValue()) + if (Value padValue = op.getPadValue()) { rewriter.create(loc, padValue); + } Value effectiveNBurst = (dmaArch == DmaArch::A5) ? op.getNBurst() : one; @@ -1199,8 +1265,9 @@ struct ExpandDmaLoadPattern : public OpRewritePattern { loc, source, destination, zero, effectiveNBurst, op.getLenBurst(), leftPadding, rightPadding, dataSelect, op.getL2CacheCtl(), op.getNburstSrcStride(), op.getNburstDstStride()); - if (hasPad) + if (hasPad) { copyOp->setAttr("has_pad", UnitAttr::get(copyOp->getContext())); + } }); if (dmaArch == DmaArch::A5 && shouldRestoreDmaLoopSize(loop1Count, loop2Size)) @@ -1405,10 +1472,12 @@ struct ExpandBiasLoadPattern : public OpRewritePattern { Value destination = materializeBufferPointer(op.getDestination(), rewriter, loc); auto sourceType = dyn_cast_or_null(source.getType()); - if (!sourceType) + if (!sourceType) { return rewriter.notifyMatchFailure(op, "expected pointer-like source"); - if (!destination) + } + if (!destination) { return rewriter.notifyMatchFailure(op, "expected pointer-like destination"); + } Value convControl = rewriter.create( loc, sourceType.getElementType().isF16() ? 1 : 0, 1); @@ -1428,8 +1497,9 @@ struct ExpandFpLoadPattern : public OpRewritePattern { Value source = materializeBufferPointer(op.getSource(), rewriter, loc); Value destination = materializeBufferPointer(op.getDestination(), rewriter, loc); - if (!source || !destination) + if (!source || !destination) { return rewriter.notifyMatchFailure(op, "expected pointer-like operands"); + } rewriter.replaceOpWithNewOp( op, source, destination, op.getNBurst(), @@ -1483,11 +1553,13 @@ struct ExpandLeftLoadPattern : public OpRewritePattern { Value destination = materializeBufferPointer(op.getDestination(), rewriter, loc); auto sourceType = dyn_cast_or_null(source.getType()); - if (!sourceType) + if (!sourceType) { return rewriter.notifyMatchFailure(op, "expected typed L1 source"); + } Type elementType = sourceType.getElementType(); - if (!destination) + if (!destination) { return rewriter.notifyMatchFailure(op, "expected pointer-like destination"); + } FailureOr control = deriveLoadCbufToCaControl( loc, op.getM(), op.getK(), elementType, op.getStartRow(), op.getStartCol(), op.getTranspose(), rewriter); @@ -1522,11 +1594,13 @@ struct ExpandRightLoadPattern : public OpRewritePattern { Value destination = materializeBufferPointer(op.getDestination(), rewriter, loc); auto sourceType = dyn_cast_or_null(source.getType()); - if (!sourceType) + if (!sourceType) { return rewriter.notifyMatchFailure(op, "expected typed L1 source"); + } Type elementType = sourceType.getElementType(); - if (!destination) + if (!destination) { return rewriter.notifyMatchFailure(op, "expected pointer-like destination"); + } FailureOr control = deriveLoadCbufToCbControl( loc, op.getK(), op.getN(), elementType, op.getStartRow(), op.getStartCol(), op.getTranspose(), rewriter); @@ -1561,10 +1635,12 @@ struct ExpandLeftLoadMxPattern : public OpRewritePattern { Value destination = materializeBufferPointer(op.getDestination(), rewriter, loc); auto sourceType = dyn_cast_or_null(source.getType()); - if (!sourceType) + if (!sourceType) { return rewriter.notifyMatchFailure(op, "expected typed L1 source"); - if (!destination) + } + if (!destination) { return rewriter.notifyMatchFailure(op, "expected pointer-like destination"); + } destination = deriveMxScaleDestination(destination, rewriter, loc); if (!destination) return rewriter.notifyMatchFailure( @@ -1610,10 +1686,12 @@ struct ExpandRightLoadMxPattern : public OpRewritePattern { Value destination = materializeBufferPointer(op.getDestination(), rewriter, loc); auto sourceType = dyn_cast_or_null(source.getType()); - if (!sourceType) + if (!sourceType) { return rewriter.notifyMatchFailure(op, "expected typed L1 source"); - if (!destination) + } + if (!destination) { return rewriter.notifyMatchFailure(op, "expected pointer-like destination"); + } destination = deriveMxScaleDestination(destination, rewriter, loc); if (!destination) return rewriter.notifyMatchFailure( @@ -1658,8 +1736,9 @@ struct ExpandAccStorePattern : public OpRewritePattern { Value source = materializeBufferPointer(op.getSource(), rewriter, loc); Value destination = materializeBufferPointer(op.getDestination(), rewriter, loc); - if (!source || !destination) + if (!source || !destination) { return rewriter.notifyMatchFailure(op, "expected pointer-like operands"); + } Value zero = getI64Constant(loc, rewriter, 0); Value one = getI64Constant(loc, rewriter, 1); configureAccStoreScalarPreOps(loc, op.getPreQuant(), op.getPreQuantMode(), @@ -1739,8 +1818,9 @@ struct ExpandAccStorePattern : public OpRewritePattern { rewriter); rewriter.create(loc, source, destination, xm, xt); - if (originalCtrl) + if (originalCtrl) { rewriter.create(loc, originalCtrl); + } rewriter.eraseOp(op); return success(); } @@ -1755,8 +1835,9 @@ struct ExpandAccStoreGmPattern : public OpRewritePattern { Value source = materializeBufferPointer(op.getSource(), rewriter, loc); Value destination = materializeBufferPointer(op.getDestination(), rewriter, loc); - if (!source || !destination) + if (!source || !destination) { return rewriter.notifyMatchFailure(op, "expected pointer-like operands"); + } Value zero = getI64Constant(loc, rewriter, 0); Value one = getI64Constant(loc, rewriter, 1); configureAccStoreScalarPreOps(loc, op.getPreQuant(), op.getPreQuantMode(), @@ -1835,8 +1916,9 @@ struct ExpandAccStoreGmPattern : public OpRewritePattern { nz2dnEn, rewriter); rewriter.create(loc, source, destination, xm, xt); - if (originalCtrl) + if (originalCtrl) { rewriter.create(loc, originalCtrl); + } rewriter.eraseOp(op); return success(); } @@ -1851,8 +1933,9 @@ struct ExpandAccStoreUbPattern : public OpRewritePattern { Value source = materializeBufferPointer(op.getSource(), rewriter, loc); Value destination = materializeBufferPointer(op.getDestination(), rewriter, loc); - if (!source || !destination) + if (!source || !destination) { return rewriter.notifyMatchFailure(op, "expected pointer-like operands"); + } Value zero = getI64Constant(loc, rewriter, 0); Value one = getI64Constant(loc, rewriter, 1); configureAccStoreScalarPreOps(loc, op.getPreQuant(), op.getPreQuantMode(), @@ -1936,8 +2019,9 @@ struct ExpandAccStoreUbPattern : public OpRewritePattern { rewriter.create(loc, source, destination, config0, config1); - if (originalCtrl) + if (originalCtrl) { rewriter.create(loc, originalCtrl); + } rewriter.eraseOp(op); return success(); } @@ -2020,8 +2104,9 @@ struct VPTOExpandWrapperOpsPass void runOnOperation() override { func::FuncOp func = getOperation(); - if (func.isExternal()) + if (func.isExternal()) { return; + } DmaArch dmaArch = getDmaArch(func->getParentOfType()); @@ -2054,8 +2139,9 @@ struct VPTOExpandWrapperOpsPass ExpandMadSemanticPattern, ExpandMadSemanticPattern, ExpandMadSemanticPattern>(&getContext()); - if (failed(applyPatternsGreedily(func, std::move(patterns)))) + if (failed(applyPatternsGreedily(func, std::move(patterns)))) { signalPassFailure(); + } } }; diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 9011e1b2e1..d9cf37e156 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -71,15 +71,25 @@ static std::optional getElementCountFromVectorLike(Type type); static Type getLowPrecisionLLVMType(Type type, MLIRContext *context) { if (pto::isPTOHiFloat8Type(type)) + { return LLVM::LLVMHiFloat8Type::get(context); + } if (isa(type)) + { return LLVM::LLVMFloat4E1M2x2Type::get(context); + } if (isa(type)) + { return LLVM::LLVMFloat4E2M1x2Type::get(context); + } if (pto::isPTOFloat8E4M3LikeType(type)) + { return LLVM::LLVMFloat8E4M3Type::get(context); + } if (pto::isPTOFloat8E5M2LikeType(type)) + { return LLVM::LLVMFloat8E5M2Type::get(context); + } return {}; } @@ -94,11 +104,15 @@ static Type normalizePayloadTypeForLLVMLowering(Type type, Builder &builder) { return getLLVMCompatibleVectorType( {2}, LLVM::LLVMHiFloat8Type::get(builder.getContext())); if (Type lowpType = getLowPrecisionLLVMType(type, builder.getContext())) + { return lowpType; + } if (auto intType = dyn_cast(type)) { if (!intType.isSignless()) + { return builder.getIntegerType(intType.getWidth()); + } return type; } @@ -106,7 +120,9 @@ static Type normalizePayloadTypeForLLVMLowering(Type type, Builder &builder) { Type normalizedElement = normalizePayloadTypeForLLVMLowering(vecType.getElementType(), builder); if (normalizedElement == vecType.getElementType()) + { return type; + } return getLLVMCompatibleVectorType(vecType.getShape(), normalizedElement, vecType.getScalableDims()); } @@ -117,9 +133,13 @@ static Type normalizePayloadTypeForLLVMLowering(Type type, Builder &builder) { static Type normalizeGEPElementTypeForLLVMLowering(Type type, Builder &builder) { if (pto::isPTOHiFloat8x2Type(type)) + { return builder.getI16Type(); + } if (pto::isPTOLowPrecisionType(type)) + { return builder.getI8Type(); + } if (isa(type)) @@ -130,7 +150,9 @@ static Type normalizeGEPElementTypeForLLVMLowering(Type type, normalizeGEPElementTypeForLLVMLowering(vecType.getElementType(), builder); if (normalizedElement == vecType.getElementType()) + { return normalizePayloadTypeForLLVMLowering(type, builder); + } return getLLVMCompatibleVectorType(vecType.getShape(), normalizedElement, vecType.getScalableDims()); } @@ -150,7 +172,9 @@ static Type convertVPTOType(Type type, Builder &builder) { if (isa(type)) return VectorType::get({32}, builder.getI8Type()); if (isa(type)) + { return LLVM::LLVMPointerType::get(builder.getContext()); + } if (auto ptrType = dyn_cast(type)) { return LLVM::LLVMPointerType::get( builder.getContext(), @@ -163,36 +187,56 @@ static unsigned getNaturalByteAlignment(Type type) { if (auto vecType = dyn_cast(type)) { unsigned elemAlign = getNaturalByteAlignment(vecType.getElementType()); if (!elemAlign) + { return 0; + } int64_t elems = 1; for (int64_t dim : vecType.getShape()) + { elems *= dim; + } return elemAlign * static_cast(elems); } if (auto intType = dyn_cast(type)) + { return llvm::divideCeil(unsigned(intType.getWidth()), 8u); + } if (pto::isPTOHiFloat8x2Type(type)) + { return 2; + } if (pto::isPTOLowPrecisionType(type)) + { return 1; + } if (type.isF16() || type.isBF16()) + { return 2; + } if (type.isF32()) + { return 4; + } if (type.isF64()) + { return 8; + } return 0; } static bool hasVPTOConvertibleType(Type type) { if (!type) + { return false; + } if (isa(type) || pto::isPTOLowPrecisionType(type)) return true; if (auto vecType = dyn_cast(type)) + { return hasVPTOConvertibleType(vecType.getElementType()); + } return false; } @@ -274,17 +318,23 @@ getVPTOStructFieldAddress(ConversionPatternRewriter &rewriter, Location loc, pto::StructType currentType = rootType; for (auto [depth, index] : llvm::enumerate(path)) { if (index < 0 || index >= static_cast(currentType.getNumFields())) + { return failure(); + } Type storageType = getVPTOStructStorageType(currentType, rewriter); address = rewriter.create( loc, pointerType, storageType, address, ArrayRef{0, static_cast(index)}); Type fieldType = currentType.getFieldType(static_cast(index)); if (depth + 1 == path.size()) + { continue; + } auto nestedStruct = dyn_cast(fieldType); if (!nestedStruct) + { return failure(); + } currentType = nestedStruct; } return address; @@ -370,9 +420,13 @@ static Value getI32Constant(OpBuilder &builder, Location loc, uint64_t value) { static bool isMxElementType(Type ty) { if (auto floatType = dyn_cast(ty)) + { return floatType.getWidth() == 8; + } if (isa(ty)) + { return true; + } std::string typeText; llvm::raw_string_ostream os(typeText); ty.print(os); @@ -382,9 +436,13 @@ static bool isMxElementType(Type ty) { static std::string getMadMxElementFragment(Type type) { if (type.isF16()) + { return "f16"; + } if (type.isBF16()) + { return "bf16"; + } std::string typeText; llvm::raw_string_ostream os(typeText); @@ -393,15 +451,25 @@ static std::string getMadMxElementFragment(Type type) { std::string lower = StringRef(typeText).lower(); if (StringRef(lower).contains("e4m3")) + { return "e4m3"; + } if (StringRef(lower).contains("e5m2")) + { return "e5m2"; + } if (StringRef(lower).contains("hif4")) + { return "hif4"; + } if (StringRef(lower).contains("e2m1x2")) + { return "e2m1x2"; + } if (StringRef(lower).contains("e1m2x2")) + { return "e1m2x2"; + } return {}; } @@ -410,7 +478,9 @@ static FailureOr buildMadMxCalleeName(MLIRContext *context, std::string lhs = getMadMxElementFragment(lhsElem); std::string rhs = getMadMxElementFragment(rhsElem); if (lhs.empty() || rhs.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.MMAD.MX." + lhs + rhs).getValue(); } @@ -421,18 +491,30 @@ static bool isSignedOrSignlessInteger(IntegerType intType, unsigned width) { static std::string getMadRhsFragment(Type type) { if (type.isF16()) + { return "f16"; + } if (type.isBF16()) + { return "bf16"; + } if (type.isF32()) + { return "f32"; + } if (auto intType = dyn_cast(type)) { if (isSignedOrSignlessInteger(intType, 4)) + { return "s4"; + } if (isSignedOrSignlessInteger(intType, 8)) + { return "s8"; + } if (intType.isUnsigned() && intType.getWidth() == 2) + { return "u2"; + } } std::string typeText; @@ -441,7 +523,9 @@ static std::string getMadRhsFragment(Type type) { os.flush(); std::string lower = StringRef(typeText).lower(); if (StringRef(lower).contains("e8m0")) + { return "e8m0"; + } return {}; } @@ -455,12 +539,18 @@ static bool isMadE5M2ElementType(Type type) { static std::string getMadDstFragment(Type type) { if (type.isF16()) + { return "f16"; + } if (type.isF32()) + { return "f32"; + } if (auto intType = dyn_cast(type)) { if (isSignedOrSignlessInteger(intType, 32)) + { return "s32"; + } } return {}; } @@ -471,15 +561,25 @@ static FailureOr buildMadTypedCalleeName(MLIRContext *context, std::string rhs = getMadRhsFragment(rhsElem); std::string dst = getMadDstFragment(dstElem); if (lhsElem.isF16() && rhs == "f16" && dst == "f32") + { return StringAttr::get(context, "llvm.hivm.MAD.f162f32.c310").getValue(); + } if (lhsElem.isF16() && rhs == "f16" && dst == "f16") + { return StringAttr::get(context, "llvm.hivm.MAD.f162f16").getValue(); + } if (lhsElem.isF16() && rhs == "f16" && dst == "s32") + { return StringAttr::get(context, "llvm.hivm.MAD.f162s32.1952").getValue(); + } if (lhsElem.isBF16() && rhs == "bf16" && dst == "f32") + { return StringAttr::get(context, "llvm.hivm.MAD.bf162f32.c310").getValue(); + } if (lhsElem.isF32() && rhs == "f32" && dst == "f32") + { return StringAttr::get(context, "llvm.hivm.MAD.f322f32.c310").getValue(); + } if (isSignedOrSignlessInteger(dyn_cast(lhsElem), 8) && rhs == "s8" && dst == "s32") return StringAttr::get(context, "llvm.hivm.MAD.s8.c310").getValue(); @@ -499,13 +599,21 @@ static FailureOr buildMadTypedCalleeName(MLIRContext *context, dst == "f32") return StringAttr::get(context, "llvm.hivm.MAD.e4m3e4m3.c310").getValue(); if (lhsElem.isF16() && rhs == "s4") + { return StringAttr::get(context, "llvm.hivm.MAD.f16s4.c310").getValue(); + } if (lhsElem.isF16() && rhs == "s8") + { return StringAttr::get(context, "llvm.hivm.MAD.f16s8.c310").getValue(); + } if (lhsElem.isF16() && rhs == "u2") + { return StringAttr::get(context, "llvm.hivm.MAD.f16u2").getValue(); + } if (lhsElem.isF16() && rhs == "e8m0") + { return StringAttr::get(context, "llvm.hivm.MAD.f16e8m0.c310").getValue(); + } return failure(); } @@ -517,7 +625,9 @@ static FailureOr buildLaneTypedCallee(MLIRContext *context, getElementTypeFragment(getElementTypeFromVectorLike(resultType)); auto lanes = getElementCountFromVectorLike(resultType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm." + stem.str() + ".v" + std::to_string(*lanes) + vec + @@ -533,7 +643,9 @@ static FailureOr buildLaneTypedCalleeFromInput(MLIRContext *context, getElementTypeFragment(getElementTypeFromVectorLike(inputType)); auto lanes = getElementCountFromVectorLike(inputType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm." + stem.str() + ".v" + std::to_string(*lanes) + vec + @@ -543,37 +655,61 @@ static FailureOr buildLaneTypedCalleeFromInput(MLIRContext *context, static std::string getElementTypeFragment(Type type) { if (type.isF16()) + { return "f16"; + } if (type.isBF16()) + { return "bf16"; + } if (type.isF32()) + { return "f32"; + } if (auto intType = dyn_cast(type)) + { return (intType.isUnsigned() ? "u" : "s") + std::to_string(intType.getWidth()); + } return {}; } static std::string getLowPrecisionElementFragment(Type type) { if (pto::isPTOHiFloat8x2Type(type)) + { return "hif8x2"; + } if (pto::isPTOHiFloat8Type(type)) + { return "hif8"; + } if (isa(type)) + { return "f4e1m2x2"; + } if (isa(type)) + { return "f4e2m1x2"; + } if (pto::isPTOFloat8E4M3LikeType(type)) + { return "f8e4m3"; + } if (pto::isPTOFloat8E5M2LikeType(type)) + { return "f8e5m2"; + } return {}; } static std::string getMemoryElementTypeFragment(Type type) { if (auto intType = dyn_cast(type)) + { return "i" + std::to_string(intType.getWidth()); + } if (std::string elem = getElementTypeFragment(type); !elem.empty()) + { return elem; + } return getLowPrecisionElementFragment(type); } @@ -590,15 +726,21 @@ struct LowpPayloadABI { static std::optional getLowpPayloadABI(Type elementType, MLIRContext *context) { if (!isLowpPayloadElementType(elementType)) + { return std::nullopt; + } return LowpPayloadABI{IntegerType::get(context, 8), "u8"}; } static std::string getDirectLowpVLogicElementFragment(Type type) { if (pto::isPTOFloat8E4M3LikeType(type)) + { return "fp8e4m3"; + } if (pto::isPTOFloat8E5M2LikeType(type)) + { return "fp8e5m2"; + } return {}; } @@ -609,7 +751,9 @@ buildDirectLowpVLogicCallee(MLIRContext *context, Type vectorType, auto lanes = getElementCountFromVectorLike(vectorType); std::string elem = getDirectLowpVLogicElementFragment(elementType); if (elem.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm." + stem.str() + "." + mode.str() + ".v" + std::to_string(*lanes) + elem) @@ -623,7 +767,9 @@ buildLowpPayloadVLogicCallee(MLIRContext *context, Type vectorType, auto lanes = getElementCountFromVectorLike(vectorType); std::optional abi = getLowpPayloadABI(elementType, context); if (!abi || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm." + stem.str() + ".v" + std::to_string(*lanes) + abi->intrinsicElementFragment.str() + @@ -647,7 +793,9 @@ static Type getLowpPayloadCarrierType(Type vectorLikeType, static Type getPayloadABIType(Type semanticType, Type convertedType, MLIRContext *context) { if (Type carrierType = getLowpPayloadCarrierType(semanticType, context)) + { return carrierType; + } return convertedType; } @@ -657,7 +805,9 @@ static Value castToPayloadABI(Location loc, Value value, Type carrierType = getLowpPayloadCarrierType(semanticType, rewriter.getContext()); if (!carrierType || carrierType == value.getType()) + { return value; + } return rewriter.create(loc, carrierType, value); } @@ -667,7 +817,9 @@ static Value castFromPayloadABI( Type carrierType = getLowpPayloadCarrierType(semanticType, rewriter.getContext()); if (!carrierType || carrierType == convertedType) + { return value; + } return rewriter.create(loc, convertedType, value); } @@ -677,17 +829,27 @@ static std::string getAtomicElementTypeFragment(Type type, if (vecType.getRank() != 1 || vecType.getDimSize(0) != 2) return {}; if (vecType.getElementType().isF16()) + { return "f16x2"; + } if (vecType.getElementType().isBF16()) + { return "bf16x2"; + } return {}; } if (type.isF16()) + { return "fp16"; + } if (type.isBF16()) + { return "bf16"; + } if (type.isF32()) + { return "fp32"; + } auto intType = dyn_cast(type); if (!intType) return {}; @@ -705,7 +867,9 @@ static std::string getAtomicElementTypeFragment(Type type, static std::string getL0LoadElementFragment(Type type) { std::string elem = getElementTypeFragment(type); if (!elem.empty()) + { return elem; + } std::string typeText; llvm::raw_string_ostream os(typeText); @@ -724,13 +888,21 @@ static std::string getL0LoadElementFragment(Type type) { static std::string getVbrScalarFragment(Type type) { if (type.isF16()) + { return "f16"; + } if (type.isBF16()) + { return "bf16"; + } if (type.isF32()) + { return "f32"; + } if (auto intType = dyn_cast(type)) + { return (intType.isUnsigned() ? "u" : "s") + std::to_string(intType.getWidth()); + } return {}; } @@ -746,9 +918,13 @@ static std::string getShuffleIntrinsicTypeFragment(Type type) { } } if (type.isF16()) + { return "f16"; + } if (type.isF32()) + { return "f32"; + } if (auto vecType = dyn_cast(type)) { if (vecType.getRank() == 1 && vecType.getDimSize(0) == 2 && vecType.getElementType().isF16()) @@ -770,26 +946,38 @@ static std::string getReduxIntrinsicTypeFragment(Type type, return isUnsigned ? "u32" : "s32"; } if (type.isF16()) + { return "f16"; + } if (type.isF32()) + { return "f32"; + } return {}; } static Type getElementTypeFromVectorLike(Type type) { if (auto vecType = dyn_cast(type)) + { return vecType.getElementType(); + } if (auto vecType = dyn_cast(type)) + { return vecType.getElementType(); + } return {}; } static std::optional getElementCountFromVectorLike(Type type) { if (auto vecType = dyn_cast(type)) + { return vecType.getElementCount(); + } if (auto vecType = dyn_cast(type)) { if (vecType.getRank() != 1) + { return std::nullopt; + } return vecType.getShape().front(); } return std::nullopt; @@ -800,21 +988,31 @@ static Value castIntegerLikeTo(Operation *anchor, Value value, Type targetType) builder.setInsertionPoint(anchor); if (value.getType() == targetType) + { return value; + } auto targetInt = dyn_cast(targetType); if (value.getType().isIndex() && targetInt) + { return builder.create(anchor->getLoc(), targetType, value); + } if (auto sourceInt = dyn_cast(value.getType())) { if (targetInt) { if (sourceInt.getWidth() < targetInt.getWidth()) + { return builder.create(anchor->getLoc(), targetType, value); + } if (sourceInt.getWidth() > targetInt.getWidth()) + { return builder.create(anchor->getLoc(), targetType, value); + } return value; } if (targetType.isIndex()) + { return builder.create(anchor->getLoc(), targetType, value); + } } return {}; @@ -825,9 +1023,13 @@ static FailureOr reinterpretPointerToAddrSpace(Operation *anchor, unsigned targetAddressSpace) { auto sourcePtrType = dyn_cast(value.getType()); if (!sourcePtrType) + { return failure(); + } if (sourcePtrType.getAddressSpace() == targetAddressSpace) + { return value; + } OpBuilder builder(anchor); builder.setInsertionPoint(anchor); @@ -843,15 +1045,21 @@ static FailureOr normalizeVdupScalarOperand(OpBuilder &builder, Location Type resultType) { auto intType = dyn_cast(input.getType()); if (!intType || intType.getWidth() != 8) + { return input; + } Type resultElemType = getElementTypeFromVectorLike(resultType); std::string resultElemFragment = getElementTypeFragment(resultElemType); if (resultElemFragment != "s8" && resultElemFragment != "u8") + { return input; + } if (intType.isSignless()) + { return input; + } Type signlessType = builder.getIntegerType(intType.getWidth()); return builder @@ -864,29 +1072,41 @@ static Value normalizeByteScalarOperandForHivmCall(OpBuilder &builder, Location Type semanticElementType) { auto intType = dyn_cast(input.getType()); if (!intType || intType.getWidth() != 8) + { return input; + } Type i16Type = builder.getIntegerType(16); auto semanticIntType = dyn_cast(semanticElementType); if (semanticIntType && semanticIntType.isUnsigned()) + { return builder.create(loc, i16Type, input).getResult(); + } return builder.create(loc, i16Type, input).getResult(); } static bool isCompatibleScalarForSemanticType(Type semanticType, Type scalarType) { if (semanticType == scalarType) + { return true; + } auto semanticInt = dyn_cast(semanticType); auto scalarInt = dyn_cast(scalarType); if (!semanticInt || !scalarInt || semanticInt.getWidth() != scalarInt.getWidth()) + { return false; + } if (semanticInt.isSigned()) + { return scalarInt.isSigned() || scalarInt.isSignless(); + } if (semanticInt.isUnsigned()) + { return scalarInt.isUnsigned() || scalarInt.isSignless(); + } return scalarInt.isSignless(); } @@ -894,11 +1114,17 @@ static std::string getCopyElementFragment(Type elementType) { if (!elementType) return {}; if (elementType.isF16()) + { return "f16"; + } if (elementType.isBF16()) + { return "bf16"; + } if (elementType.isF32()) + { return "f32"; + } // Handle FP8 family (e4m3/e5m2/e8m0/hif8) used by cube-matmul/mad_mx. std::string typeText; llvm::raw_string_ostream os(typeText); @@ -906,15 +1132,25 @@ static std::string getCopyElementFragment(Type elementType) { os.flush(); std::string lower = StringRef(typeText).lower(); if (StringRef(lower).contains("e4m3")) + { return "e4m3"; + } if (StringRef(lower).contains("e5m2")) + { return "e5m2"; + } if (StringRef(lower).contains("e8m0")) + { return "e8m0"; + } if (StringRef(lower).contains("hif8")) + { return "hif8"; + } if (StringRef(lower).contains("e1m2x2") || StringRef(lower).contains("e2m1x2")) + { return "u8"; + } if (auto intType = dyn_cast(elementType)) { switch (intType.getWidth()) { case 8: @@ -942,12 +1178,18 @@ static std::string getNd2NzCopyElementFragment(Type elementType) { StringRef(lower).contains("e8m0") || StringRef(lower).contains("hif8")) return "U8"; if (StringRef(lower).contains("e1m2x2") || StringRef(lower).contains("e2m1x2")) + { return "U8"; + } if (elementType.isF16() || elementType.isBF16()) + { return "U16"; + } if (elementType.isF32()) + { return "U32"; + } if (auto intType = dyn_cast(elementType)) { switch (intType.getWidth()) { case 8: @@ -965,77 +1207,133 @@ static std::string getNd2NzCopyElementFragment(Type elementType) { static std::optional parsePredicatePatternImmediate(StringRef pattern) { if (pattern == "PAT_ALL") + { return 0; + } if (pattern == "PAT_VL1") + { return 1; + } if (pattern == "PAT_VL2") + { return 2; + } if (pattern == "PAT_VL3") + { return 3; + } if (pattern == "PAT_VL4") + { return 4; + } if (pattern == "PAT_VL8") + { return 5; + } if (pattern == "PAT_VL16") + { return 6; + } if (pattern == "PAT_VL32") + { return 7; + } if (pattern == "PAT_VL64") + { return 8; + } if (pattern == "PAT_VL128") + { return 9; + } if (pattern == "PAT_M3") + { return 10; + } if (pattern == "PAT_M4") + { return 11; + } if (pattern == "PAT_H") + { return 12; + } if (pattern == "PAT_Q") + { return 13; + } if (pattern == "PAT_ALLF") + { return 15; + } return std::nullopt; } static std::optional parseHiLoPartImmediate(StringRef part) { if (part == "LOWER") + { return 0; + } if (part == "HIGHER") + { return 1; + } return std::nullopt; } static std::optional parseRoundModeImmediate(StringRef roundMode) { if (roundMode == "R" || roundMode == "ROUND_R") + { return 0; + } if (roundMode == "A" || roundMode == "ROUND_A") + { return 1; + } if (roundMode == "F" || roundMode == "ROUND_F") + { return 2; + } if (roundMode == "C" || roundMode == "ROUND_C") + { return 3; + } if (roundMode == "Z" || roundMode == "ROUND_Z") + { return 4; + } if (roundMode == "O" || roundMode == "ROUND_O") + { return 5; + } if (roundMode == "H" || roundMode == "ROUND_H") + { return 6; + } return std::nullopt; } static std::optional parseSaturationImmediate(StringRef sat) { if (sat == "SAT") + { return 1; + } if (sat == "NOSAT") + { return 0; + } return std::nullopt; } static std::optional parsePartImmediate(StringRef part) { if (part == "EVEN" || part == "PART_EVEN") + { return 0; + } if (part == "ODD" || part == "PART_ODD") + { return 1; + } return std::nullopt; } @@ -1047,114 +1345,190 @@ static std::optional parseVcvtPartImmediate(StringRef part) { part == "PART_P1") return 1; if (part == "P2" || part == "PART_P2") + { return 2; + } if (part == "P3" || part == "PART_P3") + { return 3; + } return std::nullopt; } static std::optional parsePredicateStoreDistImmediate(StringRef dist) { if (dist == "NORM") + { return 0; + } if (dist == "PK") + { return 1; + } return std::nullopt; } static std::optional parsePredicateLoadDistImmediate(StringRef dist) { if (dist.empty() || dist == "NORM") + { return 0; + } if (dist == "US") + { return 1; + } if (dist == "DS") + { return 2; + } return std::nullopt; } static std::optional parsePostModeImmediate(StringRef mode) { if (mode == "NO_POST_UPDATE") + { return 0; + } if (mode == "POST_UPDATE") + { return 1; + } return std::nullopt; } static std::optional parsePipeImmediate(StringRef pipe) { if (pipe == "PIPE_S") + { return 0; + } if (pipe == "PIPE_V") + { return 1; + } if (pipe == "PIPE_M") + { return 2; + } if (pipe == "PIPE_MTE1") + { return 3; + } if (pipe == "PIPE_MTE2") + { return 4; + } if (pipe == "PIPE_MTE3") + { return 5; + } if (pipe == "PIPE_ALL") + { return 6; + } if (pipe == "PIPE_MTE4") + { return 7; + } if (pipe == "PIPE_MTE5") + { return 8; + } if (pipe == "PIPE_V2") + { return 9; + } if (pipe == "PIPE_FIX") + { return 10; + } if (pipe == "VIRTUAL_PIPE_MTE2_L1A") + { return 11; + } if (pipe == "VIRTUAL_PIPE_MTE2_L1B") + { return 12; + } return std::nullopt; } static std::optional parseEventImmediate(StringRef event) { if (!event.consume_front("EVENT_ID")) + { return std::nullopt; + } uint64_t value = 0; if (event.getAsInteger(10, value)) + { return std::nullopt; + } return value; } static std::optional parseSprImmediate(StringRef spr) { if (spr == "AR") + { return 74; + } return std::nullopt; } static std::optional getDistElementWidth(Type type) { if (auto intType = dyn_cast(type)) + { return intType.getWidth(); + } if (isLowpPayloadElementType(type)) + { return 8; + } if (type.isF16() || type.isBF16()) + { return 16; + } if (type.isF32()) + { return 32; + } if (type.isF64()) + { return 64; + } return std::nullopt; } static VcvtElemKind classifyVcvtElemType(Type type) { if (type.isF16()) + { return VcvtElemKind::F16; + } if (type.isBF16()) + { return VcvtElemKind::BF16; + } if (type.isF32()) + { return VcvtElemKind::F32; + } if (pto::isPTOFloat8E4M3LikeType(type)) + { return VcvtElemKind::F8E4M3; + } if (pto::isPTOFloat8E5M2LikeType(type)) + { return VcvtElemKind::F8E5M2; + } if (pto::isPTOHiFloat8Type(type)) + { return VcvtElemKind::HiF8; + } if (isa(type)) + { return VcvtElemKind::F4E1M2x2; + } if (isa(type)) + { return VcvtElemKind::F4E2M1x2; + } if (auto intType = dyn_cast(type)) { switch (intType.getWidth()) { case 8: @@ -1366,9 +1740,13 @@ static uint64_t determineVsqzStoreHint(pto::VsqzOp vsqz) { for (Operation *user : result.getUsers()) { auto vstur = dyn_cast(user); if (!vstur) + { continue; + } if (vstur.getValue() == result) + { return 1; + } } return 0; } @@ -1377,43 +1755,81 @@ static std::optional parseLoadDistImmediate(StringRef dist, Type elementType) { auto width = getDistElementWidth(elementType); if (dist.empty() || dist == "NORM") + { return 0; + } if (!width) + { return std::nullopt; + } if (dist == "BRC_B8") + { return std::optional(1); + } if (dist == "BRC_B16") + { return std::optional(2); + } if (dist == "BRC_B32") + { return std::optional(3); + } if (dist == "US_B8") + { return std::optional(6); + } if (dist == "US_B16") + { return std::optional(7); + } if (dist == "DS_B8") + { return std::optional(8); + } if (dist == "DS_B16") + { return std::optional(9); + } if (dist == "UNPK_B8") + { return std::optional(13); + } if (dist == "UNPK_B16") + { return std::optional(14); + } if (dist == "UNPK_B32") + { return std::optional(18); + } if (dist == "BRC_BLK") + { return 15; + } if (dist == "E2B_B16") + { return std::optional(16); + } if (dist == "E2B_B32") + { return std::optional(17); + } if (dist == "UNPK4") + { return *width == 8 ? std::optional(20) : std::nullopt; + } if (dist == "SPLT4CHN") + { return *width == 8 ? std::optional(21) : std::nullopt; + } if (dist == "SPLT2CHN_B8") + { return std::optional(22); + } if (dist == "SPLT2CHN_B16") + { return std::optional(23); + } return std::nullopt; } @@ -1421,15 +1837,25 @@ static std::optional parseLoadX2DistImmediate(StringRef dist, Type elementType) { auto width = getDistElementWidth(elementType); if (dist == "BDINTLV") + { return 10; + } if (!width) + { return std::nullopt; + } if (dist == "DINTLV_B8") + { return std::optional(11); + } if (dist == "DINTLV_B16") + { return std::optional(12); + } if (dist == "DINTLV_B32") + { return std::optional(19); + } return std::nullopt; } @@ -1438,41 +1864,75 @@ static std::optional parseStoreDistImmediate(StringRef dist, auto width = getDistElementWidth(elementType); if (dist.empty()) { if (!width) + { return std::nullopt; + } if (*width == 8) + { return 0; + } if (*width == 16) + { return 1; + } if (*width == 32) + { return 2; + } return std::nullopt; } if (dist == "NORM_B8") + { return std::optional(0); + } if (dist == "NORM_B16") + { return std::optional(1); + } if (dist == "NORM_B32") + { return std::optional(2); + } if (dist == "1PT_B8") + { return std::optional(3); + } if (dist == "1PT_B16") + { return std::optional(4); + } if (dist == "1PT_B32") + { return std::optional(5); + } if (dist == "PK_B16") + { return std::optional(6); + } if (dist == "PK_B32") + { return std::optional(7); + } if (dist == "PK_B64") + { return std::optional(10); + } if (dist == "PK4_B32") + { return std::optional(12); + } if (dist == "MRG4CHN_B8") + { return std::optional(13); + } if (dist == "MRG2CHN_B8") + { return std::optional(14); + } if (dist == "MRG2CHN_B16") + { return std::optional(15); + } return std::nullopt; } @@ -1491,13 +1951,21 @@ static std::optional parseStoreX2DistImmediate(StringRef dist, Type elementType) { auto width = getDistElementWidth(elementType); if (!width) + { return std::nullopt; + } if (dist == "INTLV_B8") + { return std::optional(8); + } if (dist == "INTLV_B16") + { return std::optional(9); + } if (dist == "INTLV_B32") + { return std::optional(11); + } return std::nullopt; } @@ -1522,9 +1990,13 @@ static Value packBlockRepeatStride(Operation *anchor, Value blockStride, static std::optional parseOrderImmediate(StringRef order) { if (order.empty() || order == "ASC") + { return 0; + } if (order == "DESC") + { return 1; + } return std::nullopt; } @@ -1535,7 +2007,9 @@ static FailureOr packLoopPair(Operation *anchor, Value low, Value high) { Value lowI64 = castIntegerLikeTo(anchor, low, builder.getI64Type()); Value highI64 = castIntegerLikeTo(anchor, high, builder.getI64Type()); if (!lowI64 || !highI64) + { return failure(); + } Value shift = getI64Constant(builder, anchor->getLoc(), 40); Value highShifted = @@ -1551,7 +2025,9 @@ static FailureOr packLoopSize(Operation *anchor, Value loop2, Value loop1 Value loop2I64 = castIntegerLikeTo(anchor, loop2, builder.getI64Type()); Value loop1I64 = castIntegerLikeTo(anchor, loop1, builder.getI64Type()); if (!loop2I64 || !loop1I64) + { return failure(); + } Value shift = getI64Constant(builder, anchor->getLoc(), 21); Value loop2Shifted = @@ -1563,7 +2039,9 @@ static FailureOr packLoopSize(Operation *anchor, Value loop2, Value loop1 static FailureOr packCopyGmToUbConfig0(Operation *anchor, ValueRange operands) { if (operands.size() != 11) + { return failure(); + } OpBuilder builder(anchor); builder.setInsertionPoint(anchor); @@ -1605,7 +2083,9 @@ packCopyGmToUbConfig0(Operation *anchor, ValueRange operands) { static FailureOr packCopyGmToUbConfig1(Operation *anchor, ValueRange operands) { if (operands.size() != 11) + { return failure(); + } return packLoopPair(anchor, operands[9], operands[10]); } @@ -1622,7 +2102,9 @@ packCopyGmToUbCfgV220(Operation *anchor, ValueRange operands) { Value sid = getI64Operand(2); Value lenBurst = getI64Operand(4); if (!sid || !lenBurst) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -1666,7 +2148,9 @@ packCopyGmToUbConfig0(Operation *anchor, Value sid, Value nBurst, static FailureOr packCopyUbToGmConfig0(Operation *anchor, ValueRange operands) { if (operands.size() != 8) + { return failure(); + } OpBuilder builder(anchor); builder.setInsertionPoint(anchor); @@ -1681,7 +2165,9 @@ packCopyUbToGmConfig0(Operation *anchor, ValueRange operands) { Value lenBurst = getI64Operand(4); Value l2CacheCtl = getI64Operand(5); if (!sid || !nBurst || !lenBurst || !l2CacheCtl) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -1701,14 +2187,18 @@ packCopyUbToGmConfig0(Operation *anchor, ValueRange operands) { static FailureOr packCopyUbToGmConfig1(Operation *anchor, ValueRange operands) { if (operands.size() != 8) + { return failure(); + } return packLoopPair(anchor, operands[6], operands[7]); } static FailureOr packCopyUbToGmCfgV220(Operation *anchor, ValueRange operands) { if (operands.size() != 8) + { return failure(); + } OpBuilder builder(anchor); builder.setInsertionPoint(anchor); @@ -1721,7 +2211,9 @@ packCopyUbToGmCfgV220(Operation *anchor, ValueRange operands) { Value sid = getI64Operand(2); Value lenBurst = getI64Operand(4); if (!sid || !lenBurst) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -1761,7 +2253,9 @@ packCopyUbToGmConfig0(Operation *anchor, Value sid, Value nBurst, static FailureOr packCopyUbToUbConfig(Operation *anchor, ValueRange operands) { if (operands.size() != 7) + { return failure(); + } OpBuilder builder(anchor); builder.setInsertionPoint(anchor); Location loc = anchor->getLoc(); @@ -1775,7 +2269,9 @@ packCopyUbToUbConfig(Operation *anchor, ValueRange operands) { Value srcStride = getI64Operand(5); Value dstStride = getI64Operand(6); if (!nBurst || !lenBurst || !srcStride || !dstStride) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -1795,7 +2291,9 @@ packCopyUbToUbConfig(Operation *anchor, ValueRange operands) { static FailureOr packCopyCbufToUbConfig(Operation *anchor, ValueRange operands) { if (operands.size() != 7) + { return failure(); + } OpBuilder builder(anchor); builder.setInsertionPoint(anchor); Location loc = anchor->getLoc(); @@ -1810,7 +2308,9 @@ packCopyCbufToUbConfig(Operation *anchor, ValueRange operands) { Value srcStride = getI64Operand(5); Value dstStride = getI64Operand(6); if (!sid || !nBurst || !lenBurst || !srcStride || !dstStride) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -1831,7 +2331,9 @@ packCopyCbufToUbConfig(Operation *anchor, ValueRange operands) { static FailureOr packCopyUbToCbufConfig(Operation *anchor, ValueRange operands) { if (operands.size() != 7) + { return failure(); + } OpBuilder builder(anchor); builder.setInsertionPoint(anchor); Location loc = anchor->getLoc(); @@ -1846,7 +2348,9 @@ packCopyUbToCbufConfig(Operation *anchor, ValueRange operands) { Value srcStride = getI64Operand(5); Value dstStride = getI64Operand(6); if (!sid || !nBurst || !lenBurst || !srcStride || !dstStride) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -1873,7 +2377,9 @@ packCopyGmToCbufConfig0(Operation *anchor, Value nBurst, Value lenBurst) { Value nBurstI64 = castIntegerLikeTo(anchor, nBurst, builder.getI64Type()); Value lenBurstI64 = castIntegerLikeTo(anchor, lenBurst, builder.getI64Type()); if (!nBurstI64 || !lenBurstI64) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -1899,7 +2405,9 @@ packCopyGmToCbufConfig1(Operation *anchor, Value srcStride, Value srcStrideI64 = castIntegerLikeTo(anchor, srcStride, builder.getI64Type()); Value dstStrideI64 = castIntegerLikeTo(anchor, dstStride, builder.getI64Type()); if (!srcStrideI64 || !dstStrideI64) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -1927,7 +2435,9 @@ packCopyGmToCbufMultiConfig0(Operation *anchor, Value sid, Value l2CacheCtlI64 = castIntegerLikeTo(anchor, l2CacheCtl, builder.getI64Type()); Value nValueI64 = castIntegerLikeTo(anchor, nValue, builder.getI64Type()); if (!sidI64 || !loop1SrcStrideI64 || !l2CacheCtlI64 || !nValueI64) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -1956,7 +2466,9 @@ packCopyGmToCbufMultiConfig1(Operation *anchor, Value dValue, castIntegerLikeTo(anchor, loop4SrcStride, builder.getI64Type()); Value smallC0EnI64 = castIntegerLikeTo(anchor, smallC0En, builder.getI64Type()); if (!dValueI64 || !loop4SrcStrideI64 || !smallC0EnI64) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -2020,7 +2532,9 @@ static FailureOr packCopyCbufToFbufConfig(Operation *anchor, Value nBurst Value sourceGapI64 = castIntegerLikeTo(anchor, sourceGap, builder.getI64Type()); Value dstGapI64 = castIntegerLikeTo(anchor, dstGap, builder.getI64Type()); if (!nBurstI64 || !lenBurstI64 || !sourceGapI64 || !dstGapI64) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -2049,7 +2563,9 @@ packLoadCbufToS4Config0(Operation *anchor, Value mStart, Value kStart, Value mStepI64 = castIntegerLikeTo(anchor, mStep, builder.getI64Type()); Value kStepI64 = castIntegerLikeTo(anchor, kStep, builder.getI64Type()); if (!mStartI64 || !kStartI64 || !mStepI64 || !kStepI64) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -2075,7 +2591,9 @@ packLoadCbufToS4Config1(Operation *anchor, Value srcStride, Value dstStride) { Value srcStrideI64 = castIntegerLikeTo(anchor, srcStride, builder.getI64Type()); Value dstStrideI64 = castIntegerLikeTo(anchor, dstStride, builder.getI64Type()); if (!srcStrideI64 || !dstStrideI64) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -2097,7 +2615,9 @@ packLoadCbufToCaConfig0(Operation *anchor, Value mStart, Value kStart, Value mStepI64 = castIntegerLikeTo(anchor, mStep, builder.getI64Type()); Value kStepI64 = castIntegerLikeTo(anchor, kStep, builder.getI64Type()); if (!mStartI64 || !kStartI64 || !mStepI64 || !kStepI64) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -2125,7 +2645,9 @@ packLoadCbufToCaConfig1(Operation *anchor, Value srcStride, Value dstStride) { Value dstStrideI64 = castIntegerLikeTo(anchor, dstStride, builder.getI64Type()); if (!srcStrideI64 || !dstStrideI64) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -2147,7 +2669,9 @@ packLoadCbufToCbConfig0(Operation *anchor, Value mStart, Value kStart, Value mStepI64 = castIntegerLikeTo(anchor, mStep, builder.getI64Type()); Value kStepI64 = castIntegerLikeTo(anchor, kStep, builder.getI64Type()); if (!mStartI64 || !kStartI64 || !mStepI64 || !kStepI64) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -2175,7 +2699,9 @@ packLoadCbufToCbConfig1(Operation *anchor, Value srcStride, Value dstStride) { Value dstStrideI64 = castIntegerLikeTo(anchor, dstStride, builder.getI64Type()); if (!srcStrideI64 || !dstStrideI64) + { return failure(); + } auto shl = [&](Value value, uint64_t amount) -> Value { return builder.create(loc, value, @@ -2208,7 +2734,9 @@ static FailureOr packVbitsortConfig(Operation *anchor, Value repeatTimes) Value repeatI64 = castIntegerLikeTo(anchor, repeatTimes, builder.getI64Type()); if (!repeatI64) + { return failure(); + } return builder .create(loc, repeatI64, getI64Constant(builder, loc, 56)) .getResult(); @@ -2221,17 +2749,23 @@ static FailureOr convertElementOffsetToBytes(Operation *anchor, Value off Value offsetI32 = castIntegerLikeTo(anchor, offset, builder.getI32Type()); if (!offsetI32) + { return failure(); + } unsigned bitWidth = 0; if (auto intType = dyn_cast(elementType)) + { bitWidth = intType.getWidth(); + } else if (isLowpPayloadElementType(elementType)) bitWidth = 8; else if (auto floatType = dyn_cast(elementType)) bitWidth = floatType.getWidth(); if (bitWidth == 0 || bitWidth % 8 != 0) + { return failure(); + } Value scale = builder.create( anchor->getLoc(), builder.getI32IntegerAttr(bitWidth / 8)); @@ -2249,7 +2783,9 @@ materializeDynamicPltMask(ConversionPatternRewriter &rewriter, laneCountI32 = castIntegerLikeTo(rewriter.getInsertionBlock()->getParentOp(), laneCountI32, i32Type); if (!laneCountI32) + { return failure(); + } } StringRef calleeName; @@ -2259,14 +2795,18 @@ materializeDynamicPltMask(ConversionPatternRewriter &rewriter, calleeName = StringRef("llvm.hivm.plt.b16.v300"); } else if (auto intType = dyn_cast(vectorElemType)) { if (intType.getWidth() == 32) + { calleeName = StringRef("llvm.hivm.plt.b32.v300"); + } else if (intType.getWidth() == 16) calleeName = StringRef("llvm.hivm.plt.b16.v300"); else if (intType.getWidth() == 8) calleeName = StringRef("llvm.hivm.plt.b8.v300"); } if (calleeName.empty()) + { return failure(); + } Type maskType = VectorType::get({256}, rewriter.getI1Type()); auto funcType = @@ -2284,7 +2824,9 @@ static FailureOr buildCarryBinaryCallee(MLIRContext *context, getElementTypeFragment(cast(resultType).getElementType()); auto lanes = getElementCountFromVectorLike(resultType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm." + stem.str() + ".v" + std::to_string(*lanes) + vec) .getValue(); @@ -2293,68 +2835,116 @@ static FailureOr buildCarryBinaryCallee(MLIRContext *context, template static StringRef getUnaryMaskedStem() { if constexpr (std::is_same_v) + { return "vabs"; + } if constexpr (std::is_same_v) + { return "vexp"; + } if constexpr (std::is_same_v) + { return "vln"; + } if constexpr (std::is_same_v) + { return "vneg"; + } if constexpr (std::is_same_v) + { return "vsqrt"; + } if constexpr (std::is_same_v) + { return "vrelu"; + } if constexpr (std::is_same_v) + { return "vnot"; + } return {}; } template static StringRef getBinaryMaskedStem() { if constexpr (std::is_same_v) + { return "vadd"; + } if constexpr (std::is_same_v) + { return "vsub"; + } if constexpr (std::is_same_v) + { return "vmul"; + } if constexpr (std::is_same_v) + { return "vdiv"; + } if constexpr (std::is_same_v) + { return "vmax"; + } if constexpr (std::is_same_v) + { return "vmin"; + } if constexpr (std::is_same_v) + { return "vand"; + } if constexpr (std::is_same_v) + { return "vor"; + } if constexpr (std::is_same_v) + { return "vxor"; + } if constexpr (std::is_same_v) + { return "vshl"; + } if constexpr (std::is_same_v) + { return "vshr"; + } if constexpr (std::is_same_v) + { return "vprelu"; + } return {}; } template static StringRef getTernaryMaskedStem() { if constexpr (std::is_same_v) + { return "vmadd"; + } return {}; } template static StringRef getCarryBinaryStem() { if constexpr (std::is_same_v) + { return "vaddc"; + } if constexpr (std::is_same_v) + { return "vsubc"; + } if constexpr (std::is_same_v) + { return "vaddcs"; + } if constexpr (std::is_same_v) + { return "vsubcs"; + } return {}; } @@ -2370,7 +2960,9 @@ static FailureOr buildVselCallee(MLIRContext *context, getElementTypeFragment(cast(resultType).getElementType()); auto lanes = getElementCountFromVectorLike(resultType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vsel.v" + std::to_string(*lanes) + vec) .getValue(); @@ -2381,7 +2973,9 @@ static FailureOr buildVselrCallee(MLIRContext *context, Type elemType = getElementTypeFromVectorLike(resultType); auto lanes = getElementCountFromVectorLike(resultType); if (!elemType || !lanes) + { return failure(); + } std::string vec = getElementTypeFragment(elemType); if (auto floatType = dyn_cast(elemType); @@ -2391,7 +2985,9 @@ static FailureOr buildVselrCallee(MLIRContext *context, getLowpPayloadABI(elemType, context)) vec = abi->intrinsicElementFragment.str(); if (vec.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vselr.v" + std::to_string(*lanes) + vec) @@ -2404,7 +3000,9 @@ static FailureOr buildVdupCallee(MLIRContext *context, pto::VdupOp op std::string vec = getElementTypeFragment(getElementTypeFromVectorLike(resultType)); auto lanes = getElementCountFromVectorLike(resultType); if (vec.empty() || !lanes) + { return failure(); + } if (isa(inputType)) { StringRef position = op.getPosition().value_or("LOWEST"); @@ -2423,16 +3021,22 @@ static FailureOr buildVbrCallee(MLIRContext *context, Type semanticElementType) { std::string scalar = getVbrScalarFragment(semanticElementType); if (scalar.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vbr." + scalar + ".v300").getValue(); } static FailureOr buildPstuCallee(MLIRContext *context, pto::PstuOp op) { if (auto maskType = dyn_cast(op.getValue().getType())) { if (maskType.isB16()) + { return StringAttr::get(context, "llvm.hivm.pstu.b16").getValue(); + } if (maskType.isB32()) + { return StringAttr::get(context, "llvm.hivm.pstu.b32").getValue(); + } } return failure(); } @@ -2443,7 +3047,9 @@ static FailureOr buildVstusCallee(MLIRContext *context, getMemoryElementTypeFragment(getElementTypeFromVectorLike(valueType)); auto lanes = getElementCountFromVectorLike(valueType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vstus.v" + std::to_string(*lanes) + vec) .getValue(); @@ -2455,7 +3061,9 @@ static FailureOr buildVstusPostCallee(MLIRContext *context, getMemoryElementTypeFragment(getElementTypeFromVectorLike(valueType)); auto lanes = getElementCountFromVectorLike(valueType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vstus.post.v" + std::to_string(*lanes) + vec) .getValue(); @@ -2698,7 +3306,9 @@ FailureOr buildShuffleCallee(MLIRContext *context, Type valueType) { std::string elem = getShuffleIntrinsicTypeFragment(valueType); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.shfl.idx." + elem).getValue(); } @@ -2707,7 +3317,9 @@ FailureOr buildShuffleCallee(MLIRContext *context, Type valueType) { std::string elem = getShuffleIntrinsicTypeFragment(valueType); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.shfl.up." + elem).getValue(); } @@ -2716,7 +3328,9 @@ FailureOr buildShuffleCallee(MLIRContext *context Type valueType) { std::string elem = getShuffleIntrinsicTypeFragment(valueType); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.shfl.down." + elem).getValue(); } @@ -2725,7 +3339,9 @@ FailureOr buildShuffleCallee(MLIRContext *context Type valueType) { std::string elem = getShuffleIntrinsicTypeFragment(valueType); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.shfl.bfly." + elem).getValue(); } @@ -2753,7 +3369,9 @@ FailureOr buildReduxCallee(MLIRContext *context, Attribute signednessAttr) { std::string elem = getReduxIntrinsicTypeFragment(valueType, signednessAttr); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.redux.add." + elem).getValue(); } @@ -2763,7 +3381,9 @@ FailureOr buildReduxCallee(MLIRContext *context, Attribute signednessAttr) { std::string elem = getReduxIntrinsicTypeFragment(valueType, signednessAttr); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.redux.max." + elem).getValue(); } @@ -2773,7 +3393,9 @@ FailureOr buildReduxCallee(MLIRContext *context, Attribute signednessAttr) { std::string elem = getReduxIntrinsicTypeFragment(valueType, signednessAttr); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.redux.min." + elem).getValue(); } @@ -2788,10 +3410,14 @@ static FailureOr buildAtomicCalleeName(MLIRContext *context, StringRef opName) { std::string elem = getAtomicElementTypeFragment(valueType, signednessAttr); if (elem.empty()) + { return failure(); + } auto ptrTy = dyn_cast(ptrType); if (!ptrTy) + { return failure(); + } StringRef space; switch (ptrTy.getMemorySpace().getAddressSpace()) { @@ -2800,7 +3426,9 @@ static FailureOr buildAtomicCalleeName(MLIRContext *context, break; case pto::AddressSpace::VEC: if (valueType.isInteger(64)) + { return failure(); + } space = "S"; break; default: @@ -2839,7 +3467,9 @@ static FailureOr buildL1CacheLoadCallee(MLIRContext *context, std::string elem; if (auto intType = dyn_cast(resultType)) { if (intType.getWidth() == 8) + { elem = "s8"; + } else if (intType.getWidth() == 16) elem = "s16"; else if (intType.getWidth() == 32) @@ -2858,14 +3488,18 @@ static FailureOr buildL1CacheLoadCallee(MLIRContext *context, } else if (pto::isPTOPackedLdgStgVectorType(resultType)) { unsigned totalBits = pto::getPTOPackedLdgStgTotalBits(resultType); if (totalBits == 16) + { elem = "s16"; + } else if (totalBits == 32) elem = "s32"; else if (totalBits == 64) elem = "s64"; } if (elem.empty()) + { return failure(); + } StringRef l1cacheName = l1cache == pto::L1Cache::Cache ? "cache" : "uncache"; return StringAttr::get(context, @@ -2879,7 +3513,9 @@ static FailureOr buildL1CacheStoreCallee(MLIRContext *context, std::string elem; if (auto intType = dyn_cast(valueType)) { if (intType.getWidth() == 8) + { elem = "b8"; + } else if (intType.getWidth() == 16) elem = "b16"; else if (intType.getWidth() == 32) @@ -2898,14 +3534,18 @@ static FailureOr buildL1CacheStoreCallee(MLIRContext *context, } else if (pto::isPTOPackedLdgStgVectorType(valueType)) { unsigned totalBits = pto::getPTOPackedLdgStgTotalBits(valueType); if (totalBits == 16) + { elem = "b16"; + } else if (totalBits == 32) elem = "b32"; else if (totalBits == 64) elem = "b64"; } if (elem.empty()) + { return failure(); + } StringRef l1cacheName = l1cache == pto::L1Cache::Cache ? "cache" : "uncache"; return StringAttr::get(context, @@ -2932,7 +3572,9 @@ buildMulhiCallee(MLIRContext *context, Type resultType, .getValue(); } if (resultType.isInteger(64) && signedness == pto::Signedness::Unsigned) + { return StringAttr::get(context, "llvm.hivm.mul64hi.ui").getValue(); + } return failure(); } @@ -2947,60 +3589,86 @@ buildMulI32ToI64Callee(MLIRContext *context, pto::Signedness signedness) { static std::string getScalarFloatBuiltinFragment(Type type) { if (type.isF32()) + { return "f32"; + } if (type.isF16()) + { return "f16"; + } if (type.isBF16()) + { return "bf16"; + } return {}; } static std::string getLLVMFloatBuiltinFragment(Type type) { std::string scalar = getScalarFloatBuiltinFragment(type); if (!scalar.empty()) + { return scalar; + } auto vecType = dyn_cast(type); if (!vecType || vecType.getRank() != 1 || vecType.getDimSize(0) != 2) return {}; Type elementType = vecType.getElementType(); if (elementType.isF16()) + { return "v2f16"; + } if (elementType.isBF16()) + { return "v2bf16"; + } return {}; } static std::string getHIVMFloatBuiltinFragment(Type type) { std::string scalar = getScalarFloatBuiltinFragment(type); if (!scalar.empty()) + { return scalar; + } auto vecType = dyn_cast(type); if (!vecType || vecType.getRank() != 1 || vecType.getDimSize(0) != 2) return {}; Type elementType = vecType.getElementType(); if (elementType.isF16()) + { return "f16x2"; + } if (elementType.isBF16()) + { return "bf16x2"; + } return {}; } static FailureOr buildSqrtCallee(MLIRContext *context, Type valueType) { std::string elem = getLLVMFloatBuiltinFragment(valueType); if (elem != "f32" && elem != "f16" && elem != "v2f16") + { return failure(); + } return StringAttr::get(context, "llvm.sqrt." + elem).getValue(); } static std::string getScalarHIVMFloatShortFragment(Type type) { if (type.isF32()) + { return "f"; + } if (type.isF16()) + { return "h"; + } if (type.isBF16()) + { return "y"; + } return {}; } @@ -3012,8 +3680,9 @@ template <> FailureOr buildUnaryScalarMathCallee(MLIRContext *context, Type valueType) { std::string elem = getLLVMFloatBuiltinFragment(valueType); - if (elem != "f16" && elem != "f32" && elem != "v2f16" && elem != "v2bf16") + if (elem != "f16" && elem != "f32" && elem != "v2f16" && elem != "v2bf16") { return failure(); + } return StringAttr::get(context, "llvm.fabs." + elem).getValue(); } @@ -3022,7 +3691,9 @@ FailureOr buildUnaryScalarMathCallee(MLIRContext *context Type valueType) { std::string elem = getLLVMFloatBuiltinFragment(valueType); if (elem != "f32" && elem != "f16" && elem != "v2f16") + { return failure(); + } return StringAttr::get(context, "llvm.exp." + elem).getValue(); } @@ -3031,7 +3702,9 @@ FailureOr buildUnaryScalarMathCallee(MLIRContext *context Type valueType) { std::string elem = getLLVMFloatBuiltinFragment(valueType); if (elem != "f32" && elem != "f16" && elem != "v2f16") + { return failure(); + } return StringAttr::get(context, "llvm.log." + elem).getValue(); } @@ -3040,7 +3713,9 @@ FailureOr buildUnaryScalarMathCallee(MLIRContext *contex Type valueType) { std::string elem = getScalarHIVMFloatShortFragment(valueType); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.ceil." + elem).getValue(); } @@ -3049,7 +3724,9 @@ FailureOr buildUnaryScalarMathCallee(MLIRContext *conte Type valueType) { std::string elem = getScalarHIVMFloatShortFragment(valueType); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.floor." + elem).getValue(); } @@ -3058,7 +3735,9 @@ FailureOr buildUnaryScalarMathCallee(MLIRContext *contex Type valueType) { std::string elem = getScalarHIVMFloatShortFragment(valueType); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.rint." + elem).getValue(); } @@ -3067,7 +3746,9 @@ FailureOr buildUnaryScalarMathCallee(MLIRContext *conte Type valueType) { std::string elem = getScalarHIVMFloatShortFragment(valueType); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.round." + elem).getValue(); } @@ -3100,14 +3781,18 @@ FailureOr buildBinaryScalarMathCallee(MLIRContext *contex Type valueType) { std::string elem = getLLVMFloatBuiltinFragment(valueType); if (elem != "f32" && elem != "f16" && elem != "v2f16") + { return failure(); + } return StringAttr::get(context, "llvm.pow." + elem).getValue(); } static FailureOr buildFmaCallee(MLIRContext *context, Type valueType) { std::string elem = getHIVMFloatBuiltinFragment(valueType); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.ffma." + elem + ".rrr").getValue(); } @@ -3121,21 +3806,35 @@ static std::string getConvertScalarFragment(Type type, !elem.empty() && !pto::isPTOFloat4PackedType(elementType)) return elem + "x2"; if (elementType.isF32()) + { return "f32x2"; + } if (elementType.isF16()) + { return "f16x2"; + } if (elementType.isBF16()) + { return "bf16x2"; + } return {}; } if (type.isF32()) + { return "fp32"; + } if (type.isF16()) + { return "fp16"; + } if (type.isBF16()) + { return "bf16"; + } if (std::string elem = getLowPrecisionElementFragment(type); !elem.empty()) + { return elem; + } auto intType = dyn_cast(type); if (!intType || (intType.getWidth() != 32 && intType.getWidth() != 64) || !signednessAttr) @@ -3151,7 +3850,9 @@ static FailureOr buildConvertCallee(MLIRContext *context, std::string src = getConvertScalarFragment(srcType, signednessAttr); std::string dst = getConvertScalarFragment(dstType, signednessAttr); if (src.empty() || dst.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm." + src + ".to." + dst) .getValue(); @@ -3163,7 +3864,9 @@ static FailureOr buildVldsPostCallee(MLIRContext *context, getMemoryElementTypeFragment(getElementTypeFromVectorLike(resultType)); auto lanes = getElementCountFromVectorLike(resultType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vldsx1.post.v" + std::to_string(*lanes) + vec) .getValue(); @@ -3175,7 +3878,9 @@ static FailureOr buildVstsPostCallee(MLIRContext *context, getMemoryElementTypeFragment(getElementTypeFromVectorLike(valueType)); auto lanes = getElementCountFromVectorLike(valueType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vstsx1.post.v" + std::to_string(*lanes) + vec) .getValue(); @@ -3191,7 +3896,9 @@ static FailureOr buildVldusCallee(MLIRContext *context, getMemoryElementTypeFragment(getElementTypeFromVectorLike(resultType)); auto lanes = getElementCountFromVectorLike(resultType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vldus.v" + std::to_string(*lanes) + vec) .getValue(); @@ -3203,7 +3910,9 @@ static FailureOr buildVldusPostCallee(MLIRContext *context, getMemoryElementTypeFragment(getElementTypeFromVectorLike(resultType)); auto lanes = getElementCountFromVectorLike(resultType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vldus.post.v" + std::to_string(*lanes) + vec) .getValue(); @@ -3214,7 +3923,9 @@ static FailureOr buildVcmpCallee(MLIRContext *context, Type inputType bool isScalarCompare) { std::string elem = getElementTypeFragment(getElementTypeFromVectorLike(inputType)); if (elem.empty()) + { return failure(); + } StringRef stem = isScalarCompare ? "vcmps" : "vcmp"; return StringAttr::get(context, "llvm.hivm." + stem.str() + "." + cmpMode.str() + "." + elem + ".z") @@ -3224,56 +3935,92 @@ static FailureOr buildVcmpCallee(MLIRContext *context, Type inputType template static StringRef getVecScalarMaskedStem() { if constexpr (std::is_same_v) + { return "vmuls"; + } if constexpr (std::is_same_v) + { return "vadds"; + } if constexpr (std::is_same_v) + { return "vmaxs"; + } if constexpr (std::is_same_v) + { return "vmins"; + } if constexpr (std::is_same_v) + { return "vlrelu"; + } if constexpr (std::is_same_v) + { return "vshls"; + } if constexpr (std::is_same_v) + { return "vshrs"; + } return {}; } template static StringRef getReductionUnaryStem() { if constexpr (std::is_same_v) + { return "vcadd"; + } if constexpr (std::is_same_v) + { return "vcmax"; + } if constexpr (std::is_same_v) + { return "vcmin"; + } if constexpr (std::is_same_v) + { return "vcgadd"; + } if constexpr (std::is_same_v) + { return "vcgmax"; + } if constexpr (std::is_same_v) + { return "vcgmin"; + } if constexpr (std::is_same_v) + { return "vcpadd"; + } return {}; } template static StringRef getHistogramCallee(MLIRContext *context) { if constexpr (std::is_same_v) + { return StringAttr::get(context, "llvm.hivm.chistv2.m").getValue(); + } if constexpr (std::is_same_v) + { return StringAttr::get(context, "llvm.hivm.dhistv2.m").getValue(); + } return {}; } template static StringRef getExtremaPredicateStem() { if constexpr (std::is_same_v) + { return "vcbmax"; + } if constexpr (std::is_same_v) + { return "vcbmin"; + } return {}; } @@ -3290,7 +4037,9 @@ static FailureOr buildCopyGmToUbCallee(MLIRContext *context, bool hasPadding) { auto ptrType = dyn_cast(sourceType); if (!ptrType) + { return failure(); + } Type elementType = ptrType.getElementType(); auto getElementSuffix = [&]() -> std::string { @@ -3305,7 +4054,9 @@ static FailureOr buildCopyGmToUbCallee(MLIRContext *context, if (hasPadding) { std::string elem = getElementSuffix(); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.MOV.OUT.TO.UB.ALIGN.V2." + elem) .getValue(); @@ -3315,7 +4066,9 @@ static FailureOr buildCopyGmToUbCallee(MLIRContext *context, std::string elem = getElementSuffix(); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.MOV.OUT.TO.UB.ALIGN.V2." + elem + ".DV") .getValue(); @@ -3348,7 +4101,9 @@ static FailureOr buildOrdinaryMadCallee(MLIRContext *context, auto rhsType = dyn_cast(op.getRhs().getType()); auto dstType = dyn_cast(op.getDst().getType()); if (!lhsType || !rhsType || !dstType) + { return failure(); + } return buildMadTypedCalleeName(context, lhsType.getElementType(), rhsType.getElementType(), @@ -3360,7 +4115,9 @@ static FailureOr buildMxMadCallee(MLIRContext *context, auto lhsType = dyn_cast(op.getLhs().getType()); auto rhsType = dyn_cast(op.getRhs().getType()); if (!lhsType || !rhsType) + { return failure(); + } if (isMxElementType(lhsType.getElementType()) && isMxElementType(rhsType.getElementType())) { return buildMadMxCalleeName(context, lhsType.getElementType(), @@ -3373,10 +4130,14 @@ static FailureOr buildCopyGmToCbufCallee(MLIRContext *context, Type sourceType) { auto ptrType = dyn_cast(sourceType); if (!ptrType) + { return failure(); + } std::string elem = getCopyElementFragment(ptrType.getElementType()); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.MOV.OUT.TO.L1.ALIGN.V2." + elem + ".DV") .getValue(); @@ -3386,10 +4147,14 @@ static FailureOr buildCopyGmToCbufMultiNd2NzCallee(MLIRContext *context, Type sourceType) { auto ptrType = dyn_cast(sourceType); if (!ptrType) + { return failure(); + } std::string elem = getNd2NzCopyElementFragment(ptrType.getElementType()); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.MOV.OUT.TO.L1.MULTI.ND2NZ." + elem + ".V310") .getValue(); @@ -3411,9 +4176,13 @@ static std::string getDn2NzCopyElementFragment(Type type) { return "u8"; if (elementType.isF16() || elementType.isBF16()) + { return "u16"; + } if (elementType.isF32()) + { return "u32"; + } if (auto intType = dyn_cast(elementType)) { switch (intType.getWidth()) { @@ -3434,10 +4203,14 @@ static FailureOr buildCopyGmToCbufMultiDn2NzCallee(MLIRContext *context, Type sourceType) { auto ptrType = dyn_cast(sourceType); if (!ptrType) + { return failure(); + } std::string elem = getDn2NzCopyElementFragment(sourceType); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.MOV.OUT.TO.L1.MULTI.DN2NZ." + elem) .getValue(); @@ -3447,10 +4220,14 @@ static FailureOr buildLoadCbufToCaCallee(MLIRContext *context, Type sourceType) { auto ptrType = dyn_cast(sourceType); if (!ptrType) + { return failure(); + } std::string elem = getL0LoadElementFragment(ptrType.getElementType()); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.LOAD.L1.TO.L0A.2Dv2." + elem) .getValue(); } @@ -3459,10 +4236,14 @@ static FailureOr buildLoadCbufToCbCallee(MLIRContext *context, Type sourceType) { auto ptrType = dyn_cast(sourceType); if (!ptrType) + { return failure(); + } std::string elem = getL0LoadElementFragment(ptrType.getElementType()); if (elem.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.LOAD.L1.TO.L0B.2Dv2." + elem) .getValue(); } @@ -3471,10 +4252,14 @@ static FailureOr buildLoadCbufToCaS4Callee(MLIRContext *context, Type sourceType) { auto ptrType = dyn_cast(sourceType); if (!ptrType) + { return failure(); + } Type elementType = ptrType.getElementType(); if (!isa(elementType)) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.LOAD.L1.TO.L0A.2Dv2.s4") .getValue(); } @@ -3483,10 +4268,14 @@ static FailureOr buildLoadCbufToCbS4Callee(MLIRContext *context, Type sourceType) { auto ptrType = dyn_cast(sourceType); if (!ptrType) + { return failure(); + } Type elementType = ptrType.getElementType(); if (!isa(elementType)) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.LOAD.L1.TO.L0B.2Dv2.s4") .getValue(); } @@ -3516,7 +4305,9 @@ static FailureOr buildCopyMatrixCcToUbCallee(MLIRContext *context, Type destinationType) { auto ptrType = dyn_cast(destinationType); if (!ptrType) + { return failure(); + } Type dstElem = ptrType.getElementType(); if (dstElem.isF16()) return StringAttr::get(context, "llvm.hivm.FIX.L0C.TO.UB.f322f16.EXT") @@ -3530,7 +4321,9 @@ static FailureOr buildCopyMatrixCcToUbCallee(MLIRContext *context, static FailureOr buildCopyCbufToBtCallee(pto::CopyCbufToBtOp op) { auto ptrType = dyn_cast(op.getSource().getType()); if (!ptrType) + { return failure(); + } Type srcElem = ptrType.getElementType(); if (srcElem.isF16()) return StringAttr::get(op.getContext(), "llvm.hivm.MOV.L1.TO.BT.f16") @@ -3657,7 +4450,9 @@ static FailureOr buildUnpackCallee(MLIRContext *context, std::string result = getElementTypeFragment(getElementTypeFromVectorLike(resultType)); if (input.empty() || result.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm." + stem.str() + "." + input + "2" + result) .getValue(); @@ -3670,7 +4465,9 @@ static FailureOr buildVpackCallee(MLIRContext *context, Type inputTyp std::string result = getElementTypeFragment(getElementTypeFromVectorLike(resultType)); if (input.empty() || result.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vpack." + input + "2" + result + ".x") .getValue(); @@ -3842,7 +4639,9 @@ static FailureOr buildVldsCallee(MLIRContext *context, Type resultTyp getMemoryElementTypeFragment(getElementTypeFromVectorLike(resultType)); auto lanes = getElementCountFromVectorLike(resultType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vldsx1.v" + std::to_string(*lanes) + vec) .getValue(); @@ -3853,10 +4652,14 @@ static FailureOr buildVldsx2Callee(MLIRContext *context, Type elementType = getElementTypeFromVectorLike(resultType); auto lanes = getElementCountFromVectorLike(resultType); if (!elementType || !lanes) + { return failure(); + } std::string element = getMemoryElementTypeFragment(elementType); if (element.empty()) + { return failure(); + } return StringAttr::get( context, "llvm.hivm.vldsx2" + std::string(post ? ".post" : "") + ".v" + @@ -3870,17 +4673,25 @@ buildBlockStridedMemoryCallee(MLIRContext *context, Type vectorType, Type elementType = getElementTypeFromVectorLike(vectorType); auto lanes = getElementCountFromVectorLike(vectorType); if (!elementType || !lanes) + { return failure(); + } std::string element; if (auto intType = dyn_cast(elementType)) + { element = "i" + std::to_string(intType.getWidth()); + } else if (isLowpPayloadElementType(elementType)) element = "i8"; else + { element = getMemoryElementTypeFragment(elementType); + } if (element.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm." + stem.str() + @@ -3900,7 +4711,9 @@ static FailureOr buildVstsCallee(MLIRContext *context, Type valueType getMemoryElementTypeFragment(getElementTypeFromVectorLike(valueType)); auto lanes = getElementCountFromVectorLike(valueType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vstsx1.v" + std::to_string(*lanes) + vec) .getValue(); @@ -3910,11 +4723,15 @@ static FailureOr buildVstsx2Callee(MLIRContext *context, Type valueTy Type elementType = getElementTypeFromVectorLike(valueType); auto lanes = getElementCountFromVectorLike(valueType); if (!elementType || !lanes) + { return failure(); + } std::string element = getMemoryElementTypeFragment(elementType); if (element.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vstsx2.v" + std::to_string(*lanes) + element) @@ -3928,9 +4745,13 @@ static FailureOr buildVsstbCallee(MLIRContext *context, static Type getVgather2SourceElementType(Type sourceType) { if (auto ptrType = dyn_cast(sourceType)) + { return ptrType.getElementType(); + } if (auto memrefType = dyn_cast(sourceType)) + { return memrefType.getElementType(); + } return {}; } @@ -3941,7 +4762,9 @@ static FailureOr buildVgather2Callee(MLIRContext *context, Type resultElemType = getElementTypeFromVectorLike(resultType); auto lanes = getElementCountFromVectorLike(resultType); if (!sourceElemType || !resultElemType || !lanes) + { return failure(); + } std::string vec; int64_t intrinsicLanes = *lanes; @@ -3952,7 +4775,9 @@ static FailureOr buildVgather2Callee(MLIRContext *context, vec = getElementTypeFragment(resultElemType); } if (vec.empty()) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vgather2.v300.v" + std::to_string(intrinsicLanes) + vec) @@ -3962,13 +4787,19 @@ static FailureOr buildVgather2Callee(MLIRContext *context, static std::optional getFixedVectorBitWidth(Type type) { auto vectorType = dyn_cast(type); if (!vectorType || vectorType.getRank() != 1 || vectorType.isScalable()) + { return std::nullopt; + } int64_t lanes = vectorType.getDimSize(0); if (lanes <= 0) + { return std::nullopt; + } auto elementType = dyn_cast(vectorType.getElementType()); if (!elementType) + { return std::nullopt; + } return static_cast(lanes) * elementType.getWidth(); } @@ -3980,19 +4811,25 @@ static FailureOr getVgather2OffsetsCarrierType(PatternRewriter &rewriter, Type elementType = getElementTypeFromVectorLike(resultType); auto lanes = getElementCountFromVectorLike(resultType); if (!sourceElemType || !elementType || !lanes || *lanes <= 0) + { return failure(); + } Type carrierType = offsetsType; if (pto::getPTOStorageElemBitWidth(elementType) == 16) { if (*lanes % 2 != 0) + { return failure(); + } carrierType = VectorType::get({*lanes / 2}, rewriter.getI32Type()); } std::optional offsetsBits = getFixedVectorBitWidth(offsetsType); std::optional carrierBits = getFixedVectorBitWidth(carrierType); if (!offsetsBits || !carrierBits || *offsetsBits != *carrierBits) + { return failure(); + } return carrierType; } @@ -4028,7 +4865,9 @@ static FailureOr buildVmulscvtCallee(MLIRContext *context, auto inputLanes = getElementCountFromVectorLike(inputType); auto resultLanes = getElementCountFromVectorLike(resultType); if (!inputElemType || !resultElemType || !inputLanes || !resultLanes) + { return failure(); + } if (!inputElemType.isF32() || !resultElemType.isF16() || *inputLanes != 64 || *resultLanes != 128) return failure(); @@ -4040,7 +4879,9 @@ static FailureOr buildVciCallee(MLIRContext *context, Type resultType getElementTypeFragment(getElementTypeFromVectorLike(resultType)); auto lanes = getElementCountFromVectorLike(resultType); if (vec.empty() || !lanes) + { return failure(); + } if (vec == "f16" || vec == "f32") return StringAttr::get(context, "llvm.hivm.vci.v" + std::to_string(*lanes) + vec + "." + vec) @@ -4055,7 +4896,9 @@ static FailureOr buildVtrcCallee(MLIRContext *context, Type resultTyp getElementTypeFragment(getElementTypeFromVectorLike(resultType)); auto lanes = getElementCountFromVectorLike(resultType); if (vec.empty() || !lanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vtrc." + vec + ".x").getValue(); } @@ -4068,7 +4911,9 @@ static FailureOr buildVexpdifCallee(MLIRContext *context, std::string dstElem = getElementTypeFragment(getElementTypeFromVectorLike(resultType)); if (srcVec.empty() || dstElem.empty() || !srcLanes) + { return failure(); + } return StringAttr::get(context, "llvm.hivm.vexpdif.v" + std::to_string(*srcLanes) + srcVec + dstElem) @@ -4079,9 +4924,13 @@ static FailureOr buildVbitsortCallee(MLIRContext *context, pto::VbitsortOp op) { Type sourceElemType = cast(op.getSource().getType()).getElementType(); if (sourceElemType.isF16()) + { return StringAttr::get(context, "llvm.hivm.VBS32.V300.f16").getValue(); + } if (sourceElemType.isF32()) + { return StringAttr::get(context, "llvm.hivm.VBS32.V300.f32").getValue(); + } return failure(); } @@ -4090,9 +4939,13 @@ static FailureOr buildVmrgsort4Callee(MLIRContext *context, Type elemType = cast(op.getDestination().getType()).getElementType(); if (elemType.isF16()) + { return StringAttr::get(context, "llvm.hivm.VMRGSORT.f16.V300").getValue(); + } if (elemType.isF32()) + { return StringAttr::get(context, "llvm.hivm.VMRGSORT.f32.V300").getValue(); + } return failure(); } @@ -4104,16 +4957,22 @@ static FailureOr packVmrgsort4SourceAddr(Operation *anchor, Value source0 Location loc = anchor->getLoc(); unsigned addrShift = 0; if (elemType.isF16()) + { addrShift = 3; + } else if (elemType.isF32()) addrShift = 3; else + { return failure(); + } auto packOne = [&](Value source, uint64_t laneShift) -> FailureOr { FailureOr ubPtr = reinterpretPointerToAddrSpace(anchor, source, 6); if (failed(ubPtr)) + { return failure(); + } Value asInt = builder.create(loc, builder.getI64Type(), *ubPtr); Value shifted = builder.create( @@ -4121,7 +4980,9 @@ static FailureOr packVmrgsort4SourceAddr(Operation *anchor, Value source0 Value masked = builder.create( loc, shifted, getI64Constant(builder, loc, 0xFFFFULL)); if (laneShift == 0) + { return masked; + } return builder .create(loc, masked, getI64Constant(builder, loc, laneShift)) @@ -4133,7 +4994,9 @@ static FailureOr packVmrgsort4SourceAddr(Operation *anchor, Value source0 FailureOr low2 = packOne(source2, 32); FailureOr low3 = packOne(source3, 48); if (failed(low0) || failed(low1) || failed(low2) || failed(low3)) + { return failure(); + } Value packed01 = builder.create(loc, *low0, *low1); Value packed23 = builder.create(loc, *low2, *low3); @@ -4146,17 +5009,23 @@ static FailureOr buildVcvtContract(pto::VcvtOp op) { Type inputElemType = getElementTypeFromVectorLike(op.getInput().getType()); Type resultElemType = getElementTypeFromVectorLike(op.getResult().getType()); if (!inputElemType || !resultElemType) + { return failure(); + } auto contract = lookupVcvtContract(classifyVcvtElemType(inputElemType), classifyVcvtElemType(resultElemType)); if (!contract) + { return failure(); + } return *contract; } static bool needsV300CtrlModeForVPTOFunc(func::FuncOp funcOp) { if (!pto::isPTOEntryFunction(funcOp) || funcOp.getBlocks().empty()) + { return false; + } bool needsCtrlSetup = false; funcOp.walk([&](pto::VcvtOp vcvtOp) { @@ -4316,7 +5185,9 @@ static FailureOr encodeMovPadValue(Location loc, Value value, } if (bitWidth != 8 && bitWidth != 16 && bitWidth != 32) + { return failure(); + } return rewriter.create(loc, rewriter.getI64Type(), payload) .getResult(); @@ -4516,7 +5387,9 @@ class LowerUnaryMaskedOpPattern final : public OpConversionPattern { FailureOr calleeName = buildLaneTypedCallee(op.getContext(), op.getResult().getType(), stem, ".x"); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported unary VPTO signature"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); @@ -4559,13 +5432,17 @@ class LowerVsqzOpPattern final : public OpConversionPattern { FailureOr calleeName = buildVsqzCallee(op.getContext(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vsqz VPTO signature"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); Type maskType = this->getTypeConverter()->convertType(op.getMask().getType()); if (!resultType || !maskType) + { return rewriter.notifyMatchFailure(op, "failed to convert vsqz types"); + } Value input = adaptor.getInput(); Value mask = adaptor.getMask(); @@ -4603,13 +5480,17 @@ class LowerVusqzOpPattern final : public OpConversionPattern { FailureOr calleeName = buildVusqzCallee(op.getContext(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vusqz VPTO signature"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); Type maskType = this->getTypeConverter()->convertType(op.getMask().getType()); if (!resultType || !maskType) + { return rewriter.notifyMatchFailure(op, "failed to convert vusqz types"); + } Value src = adaptor.getSrc(); Value mask = adaptor.getMask(); @@ -4643,13 +5524,17 @@ class LowerVmulaOpPattern final : public OpConversionPattern { FailureOr calleeName = buildVmulaCallee(op.getContext(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vmula VPTO signature"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); Type maskType = this->getTypeConverter()->convertType(op.getMask().getType()); if (!resultType || !maskType) + { return rewriter.notifyMatchFailure(op, "failed to convert vmula types"); + } Value acc = adaptor.getAcc(); Value lhs = adaptor.getLhs(); @@ -4689,7 +5574,9 @@ class LowerVmullOpPattern final : public OpConversionPattern { FailureOr calleeName = buildVmullCallee(op.getContext(), op.getLow().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vmull VPTO signature"); + } Type inputType = this->getTypeConverter()->convertType(op.getLhs().getType()); Type maskType = this->getTypeConverter()->convertType(op.getMask().getType()); @@ -4699,7 +5586,9 @@ class LowerVmullOpPattern final : public OpConversionPattern { return rewriter.notifyMatchFailure(op, "failed to convert vmull types"); } if (resultTypes.size() != 2 || resultTypes[0] != resultTypes[1]) + { return rewriter.notifyMatchFailure(op, "unexpected converted vmull results"); + } Value lhs = adaptor.getLhs(); Value rhs = adaptor.getRhs(); @@ -4782,7 +5671,9 @@ class LowerBinaryMaskedOpPattern final : public OpConversionPattern { } } if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported binary VPTO signature"); + } auto call = rewriter.create(op.getLoc(), *calleeName, TypeRange{callResultType}, @@ -4865,7 +5756,9 @@ class LowerCarryBinaryOpPattern final : public OpConversionPattern { FailureOr calleeName = buildCarryBinaryCallee(op.getContext(), op.getResult().getType(), stem); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported carry VPTO signature"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); @@ -4915,23 +5808,31 @@ class LowerCopyOpPattern final : public OpConversionPattern { bool hasPadding = false; if constexpr (isGmUb) + { hasPadding = op->hasAttr("has_pad"); + } FailureOr calleeName = failure(); if constexpr (isGmUb) calleeName = buildCopyGmToUbCallee(op.getContext(), op.getSource().getType(), march, hasPadding); else + { calleeName = buildCopyUbToGmCallee(op.getContext(), march); + } if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported copy VPTO signature"); + } auto llvmSourceType = dyn_cast(adaptor.getOperands()[0].getType()); auto llvmDestType = dyn_cast(adaptor.getOperands()[1].getType()); if (!llvmSourceType || !llvmDestType) + { return rewriter.notifyMatchFailure(op, "expected LLVM pointer copy operands"); + } bool isC220 = march == "dav-c220-vec" || march == "dav-c220-cube"; bool useA3NonPadded = isC220 && isGmUb && !hasPadding; @@ -4941,7 +5842,9 @@ class LowerCopyOpPattern final : public OpConversionPattern { FailureOr config0 = failure(); FailureOr config1 = failure(); if (useA3NonPadded) + { config0 = packCopyGmToUbCfgV220(op, adaptor.getOperands()); + } else if (useA3UbGm) config0 = packCopyUbToGmCfgV220(op, adaptor.getOperands()); else if constexpr (isGmUb) { @@ -4952,7 +5855,9 @@ class LowerCopyOpPattern final : public OpConversionPattern { config1 = packCopyUbToGmConfig1(op, adaptor.getOperands()); } if (failed(config0) || (!useSingleConfig && failed(config1))) + { return rewriter.notifyMatchFailure(op, "failed to materialize copy config"); + } SmallVector args{adaptor.getOperands()[1], adaptor.getOperands()[0], *config0}; @@ -4996,7 +5901,9 @@ class LowerUBufBinaryOpPattern final : public OpConversionPattern { std::string calleeName; if constexpr (std::is_same_v) + { calleeName = "llvm.hivm.VADD." + elemFrag; + } else if constexpr (std::is_same_v) calleeName = "llvm.hivm.VSUB." + elemFrag; else if constexpr (std::is_same_v) @@ -5014,7 +5921,9 @@ class LowerUBufBinaryOpPattern final : public OpConversionPattern { else if constexpr (std::is_same_v) calleeName = "llvm.hivm.VADDRELU." + elemFrag; else + { return rewriter.notifyMatchFailure(op, "unsupported ubuf binary op"); + } Value dst = adaptor.getDst(); Value src0 = adaptor.getSrc0(); @@ -5093,17 +6002,23 @@ class LowerUBufShiftOpPattern final : public OpConversionPattern { op, "unsupported element type for ubuf shift op"); if (elemFrag == "s16") + { elemFrag = "u16"; + } else if (elemFrag == "s32") elemFrag = "u32"; std::string calleeName; if constexpr (std::is_same_v) + { calleeName = "llvm.hivm.VSHL." + elemFrag; + } else if constexpr (std::is_same_v) calleeName = "llvm.hivm.VSHR." + elemFrag; else + { return rewriter.notifyMatchFailure(op, "unsupported ubuf shift op"); + } Value dst = adaptor.getDst(); Value src = adaptor.getSrc(); @@ -5196,7 +6111,9 @@ class LowerUBufScalarBinaryPattern final : public OpConversionPattern // Scalar-tile ops keep signed intrinsic names (s16/s32). std::string calleeName; if constexpr (std::is_same_v) + { calleeName = "llvm.hivm.VMULS." + elemFrag; + } else if constexpr (std::is_same_v) calleeName = "llvm.hivm.VADDS." + elemFrag; else if constexpr (std::is_same_v) @@ -5204,7 +6121,9 @@ class LowerUBufScalarBinaryPattern final : public OpConversionPattern else if constexpr (std::is_same_v) calleeName = "llvm.hivm.VMINS." + elemFrag; else + { return rewriter.notifyMatchFailure(op, "unsupported ubuf scalar binary op"); + } Value dst = adaptor.getDst(); Value src = adaptor.getSrc(); @@ -5292,15 +6211,21 @@ class LowerUBufVdupPattern final : public OpConversionPattern { Type elemType = ptrType.getElementType(); std::string suffix; if (elemType.isF32() || elemType.isInteger(32)) + { suffix = "u32"; + } else if (elemType.isF16() || elemType.isInteger(16)) suffix = "u16"; else + { return rewriter.notifyMatchFailure(op, "unsupported element type for ubuf vdup"); + } Value dst = adaptor.getDst(); if (!dst || !isa(dst.getType())) + { return rewriter.notifyMatchFailure(op, "unexpected converted ubuf vdup dst type"); + } Location loc = op.getLoc(); auto i64Ty = rewriter.getI64Type(); @@ -5363,11 +6288,15 @@ class LowerUBufUnaryOpPattern final : public OpConversionPattern { op, "unsupported element type for ubuf unary op"); if (elemFrag == "s16") + { elemFrag = "u16"; + } std::string calleeName; if constexpr (std::is_same_v) + { calleeName = "llvm.hivm.VNOT." + elemFrag; + } else if constexpr (std::is_same_v) calleeName = "llvm.hivm.VABS." + elemFrag; else if constexpr (std::is_same_v) { @@ -5384,7 +6313,9 @@ class LowerUBufUnaryOpPattern final : public OpConversionPattern { else if constexpr (std::is_same_v) calleeName = "llvm.hivm.VRSQRT." + elemFrag; else + { return rewriter.notifyMatchFailure(op, "unsupported ubuf unary op"); + } Value dst = adaptor.getDst(); Value src = adaptor.getSrc(); @@ -5542,11 +6473,15 @@ class LowerCopyUbufToUbufOpPattern final auto llvmDestType = dyn_cast(adaptor.getOperands()[1].getType()); if (!llvmSourceType || !llvmDestType) + { return rewriter.notifyMatchFailure(op, "expected LLVM pointer copy operands"); + } FailureOr config = packCopyUbToUbConfig(op, adaptor.getOperands()); if (failed(config)) + { return rewriter.notifyMatchFailure(op, "failed to materialize copy config"); + } StringRef calleeName = buildCopyUbToUbCallee(op.getContext()); SmallVector args{adaptor.getOperands()[1], adaptor.getOperands()[0], @@ -5582,7 +6517,9 @@ class LowerCopyCbufToUbufOpPattern final Value sourceRaw = adaptor.getSource(); Value destinationRaw = adaptor.getDestination(); if (!sourceRaw || !destinationRaw) + { return rewriter.notifyMatchFailure(op, "expected converted operands"); + } if (!isa(sourceRaw.getType()) || !isa(destinationRaw.getType())) return rewriter.notifyMatchFailure(op, "expected LLVM pointer src/dst"); @@ -5596,11 +6533,15 @@ class LowerCopyCbufToUbufOpPattern final FailureOr destination = reinterpretPointerToAddrSpace(op, destinationRaw, ubufAddressSpace); if (failed(source) || failed(destination)) + { return rewriter.notifyMatchFailure(op, "failed to map cbuf/ubuf pointer spaces"); + } FailureOr config = packCopyCbufToUbConfig(op, adaptor.getOperands()); if (failed(config)) + { return rewriter.notifyMatchFailure(op, "failed to materialize copy config"); + } StringRef calleeName = buildCopyCbufToUbCallee(op.getContext()); auto funcType = rewriter.getFunctionType( @@ -5634,7 +6575,9 @@ class LowerCopyUbufToCbufOpPattern final Value sourceRaw = adaptor.getSource(); Value destinationRaw = adaptor.getDestination(); if (!sourceRaw || !destinationRaw) + { return rewriter.notifyMatchFailure(op, "expected converted operands"); + } if (!isa(sourceRaw.getType()) || !isa(destinationRaw.getType())) return rewriter.notifyMatchFailure(op, "expected LLVM pointer src/dst"); @@ -5648,11 +6591,15 @@ class LowerCopyUbufToCbufOpPattern final FailureOr destination = reinterpretPointerToAddrSpace(op, destinationRaw, cbufAddressSpace); if (failed(source) || failed(destination)) + { return rewriter.notifyMatchFailure(op, "failed to map ubuf/cbuf pointer spaces"); + } FailureOr config = packCopyUbToCbufConfig(op, adaptor.getOperands()); if (failed(config)) + { return rewriter.notifyMatchFailure(op, "failed to materialize copy config"); + } StringRef calleeName = buildCopyUbToCbufCallee(op.getContext()); auto funcType = rewriter.getFunctionType( @@ -5708,7 +6655,9 @@ static LogicalResult lowerMadRawOp(pto::MadRawOpInterface op, reinterpretPointerToAddrSpace(op, dstRaw, ccAddressSpace); FailureOr bias; if (biasRaw) + { bias = reinterpretPointerToAddrSpace(op, biasRaw, btAddressSpace); + } if (failed(lhs) || failed(rhs) || failed(dst) || (biasRaw && failed(bias))) { return rewriter.notifyMatchFailure(op, "failed to map cube pointer spaces"); @@ -5723,7 +6672,9 @@ static LogicalResult lowerMadRawOp(pto::MadRawOpInterface op, Value callDst = *dst; if (biasRaw) + { callDst = buildMadBiasDestination(op, rewriter, *dst, *bias); + } auto funcType = rewriter.getFunctionType( TypeRange{dst->getType(), lhs->getType(), rhs->getType(), i64Ty}, TypeRange{}); @@ -5747,7 +6698,9 @@ class LowerMadRawPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto raw = dyn_cast(op.getOperation()); if (!raw) + { return failure(); + } return lowerMadRawOp(raw, adaptor.getOperands(), rewriter, state); } @@ -5798,7 +6751,9 @@ class LowerCopyGmToCbufOpPattern final FailureOr destination = reinterpretPointerToAddrSpace(op, destinationRaw, cbufAddressSpace); if (failed(source) || failed(destination)) + { return rewriter.notifyMatchFailure(op, "failed to map cbuf/gm pointer spaces"); + } FailureOr calleeName = buildCopyGmToCbufCallee(op.getContext(), op.getSource().getType()); @@ -5843,7 +6798,9 @@ class LowerCopyGmToCbufMultiOpPattern final Value sourceRaw = adaptor.getSource(); Value destinationRaw = adaptor.getDestination(); if (!sourceRaw || !destinationRaw) + { return rewriter.notifyMatchFailure(op, "expected converted operands"); + } if (!isa(sourceRaw.getType()) || !isa(destinationRaw.getType())) return rewriter.notifyMatchFailure(op, "expected LLVM pointer src/dst"); @@ -5857,7 +6814,9 @@ class LowerCopyGmToCbufMultiOpPattern final FailureOr destination = reinterpretPointerToAddrSpace(op, destinationRaw, cbufAddressSpace); if (failed(source) || failed(destination)) + { return rewriter.notifyMatchFailure(op, "failed to map cbuf/gm pointer spaces"); + } FailureOr config0 = packCopyGmToCbufMultiConfig0( op, adaptor.getSid(), adaptor.getLoop1SrcStride(), @@ -5867,12 +6826,16 @@ class LowerCopyGmToCbufMultiOpPattern final adaptor.getLoop4SrcStride(), adaptor.getSmallc0En()); if (failed(config0) || failed(config1)) + { return rewriter.notifyMatchFailure(op, "failed to pack multi copy config"); + } FailureOr calleeName = [&] (MLIRContext *ctx, Type sourceType) -> FailureOr { if constexpr (std::is_same_v) + { return buildCopyGmToCbufMultiNd2NzCallee(ctx, op.getSource().getType()); + } return buildCopyGmToCbufMultiDn2NzCallee(ctx, sourceType); }(op.getContext(), op.getSource().getType()); if (failed(calleeName)) @@ -5909,7 +6872,9 @@ class LowerCopyCbufToBtOpPattern final Value sourceRaw = adaptor.getSource(); Value destinationRaw = adaptor.getDestination(); if (!sourceRaw || !destinationRaw) + { return rewriter.notifyMatchFailure(op, "expected converted operands"); + } if (!isa(sourceRaw.getType()) || !isa(destinationRaw.getType())) return rewriter.notifyMatchFailure(op, "expected LLVM pointer src/dst"); @@ -5923,20 +6888,26 @@ class LowerCopyCbufToBtOpPattern final FailureOr destinationPtr = reinterpretPointerToAddrSpace(op, destinationRaw, btAddressSpace); if (failed(source) || failed(destinationPtr)) + { return rewriter.notifyMatchFailure(op, "failed to map cbuf/bt pointer spaces"); + } FailureOr config = packCopyCbufToBtConfig( op, adaptor.getConvControl(), adaptor.getNBurst(), adaptor.getLenBurst(), adaptor.getSourceGap(), adaptor.getDstGap()); if (failed(config)) + { return rewriter.notifyMatchFailure(op, "failed to pack copy_cbuf_to_bt config"); + } Type i64Ty = rewriter.getI64Type(); Value destination = rewriter.create(op.getLoc(), i64Ty, *destinationPtr); FailureOr calleeName = buildCopyCbufToBtCallee(op); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported copy_cbuf_to_bt source element type"); + } auto funcType = rewriter.getFunctionType( TypeRange{i64Ty, source->getType(), i64Ty}, TypeRange{}); rewriter.create(op.getLoc(), *calleeName, TypeRange{}, @@ -5965,7 +6936,9 @@ class LowerCopyCbufToFbufOpPattern final Value sourceRaw = adaptor.getSource(); Value destinationRaw = adaptor.getDestination(); if (!sourceRaw || !destinationRaw) + { return rewriter.notifyMatchFailure(op, "expected converted operands"); + } if (!isa(sourceRaw.getType()) || !isa(destinationRaw.getType())) return rewriter.notifyMatchFailure(op, "expected LLVM pointer src/dst"); @@ -5978,13 +6951,17 @@ class LowerCopyCbufToFbufOpPattern final FailureOr destination = reinterpretPointerToAddrSpace(op, destinationRaw, fbufAddressSpace); if (failed(source) || failed(destination)) + { return rewriter.notifyMatchFailure(op, "failed to map cbuf/fbuf pointer spaces"); + } FailureOr config = packCopyCbufToFbufConfig( op, adaptor.getNBurst(), adaptor.getLenBurst(), adaptor.getSourceGap(), adaptor.getDstGap()); if (failed(config)) + { return rewriter.notifyMatchFailure(op, "failed to pack copy_cbuf_to_fbuf config"); + } Type i64Ty = rewriter.getI64Type(); StringRef calleeName = buildCopyCbufToFbufCallee(op.getContext()); @@ -6040,14 +7017,18 @@ class LowerLoadCbufToCaOpPattern final FailureOr destination = reinterpretPointerToAddrSpace(op, destinationRaw, caAddressSpace); if (failed(source) || failed(destination)) + { return rewriter.notifyMatchFailure(op, "failed to map cbuf/ca pointer spaces"); + } FailureOr config0 = packLoadCbufToCaConfig0(op, mStart, kStart, mStep, kStep); FailureOr config1 = packLoadCbufToCaConfig1(op, srcStride, dstStride); if (failed(config0) || failed(config1)) + { return rewriter.notifyMatchFailure(op, "failed to pack load_cbuf_to_ca config"); + } Value transpose = getI64Constant(rewriter, op.getLoc(), op.getTranspose() ? 1 : 0); @@ -6086,7 +7067,9 @@ class LowerLoadCbufToS4OpPattern final : public OpConversionPattern { Value sourceRaw = adaptor.getSource(); Value destinationRaw = adaptor.getDestination(); if (!sourceRaw || !destinationRaw) + { return rewriter.notifyMatchFailure(op, "expected converted operands"); + } if (!isa(sourceRaw.getType()) || !isa(destinationRaw.getType())) return rewriter.notifyMatchFailure(op, "expected LLVM pointer src/dst"); @@ -6102,7 +7085,9 @@ class LowerLoadCbufToS4OpPattern final : public OpConversionPattern { FailureOr destination = reinterpretPointerToAddrSpace(op, destinationRaw, targetAddressSpace); if (failed(source) || failed(destination)) + { return rewriter.notifyMatchFailure(op, "failed to map cbuf/cube pointer spaces"); + } FailureOr config0 = packLoadCbufToS4Config0( op, adaptor.getMStart(), adaptor.getKStart(), adaptor.getMStep(), @@ -6111,12 +7096,16 @@ class LowerLoadCbufToS4OpPattern final : public OpConversionPattern { packLoadCbufToS4Config1(op, adaptor.getSrcStride(), adaptor.getDstStride()); if (failed(config0) || failed(config1)) + { return rewriter.notifyMatchFailure(op, "failed to pack load_cbuf_to_*_s4 config"); + } Value transpose = castIntegerLikeTo(op, adaptor.getTranspose(), rewriter.getI64Type()); if (!transpose) + { return rewriter.notifyMatchFailure(op, "failed to cast transpose to i64"); + } FailureOr calleeName = std::is_same_v @@ -6183,7 +7172,9 @@ class LowerLoadCbufToCbOpPattern final FailureOr destination = reinterpretPointerToAddrSpace(op, destinationRaw, cbAddressSpace); if (failed(source) || failed(destination)) + { return rewriter.notifyMatchFailure(op, "failed to map cbuf/cb pointer spaces"); + } bool transpose = op.getTranspose(); FailureOr config0 = @@ -6191,7 +7182,9 @@ class LowerLoadCbufToCbOpPattern final FailureOr config1 = packLoadCbufToCbConfig1(op, srcStride, dstStride); if (failed(config0) || failed(config1)) + { return rewriter.notifyMatchFailure(op, "failed to pack load_cbuf_to_cb config"); + } Value transposeValue = getI64Constant(rewriter, op.getLoc(), transpose ? 1 : 0); @@ -6246,7 +7239,9 @@ class LowerLoadCbufToCaMxOpPattern final FailureOr src = reinterpretPointerToAddrSpace(op, srcRaw, cbufAddressSpace); FailureOr dst = reinterpretPointerToAddrSpace(op, dstRaw, caAddressSpace); if (failed(src) || failed(dst)) + { return rewriter.notifyMatchFailure(op, "failed to map cbuf/ca pointer spaces"); + } Type sourceElemType = cast(op.getSource().getType()).getElementType(); unsigned elemBitWidth = pto::getPTOStorageElemBitWidth(sourceElemType); @@ -6311,7 +7306,9 @@ class LowerLoadCbufToCbMxOpPattern final FailureOr src = reinterpretPointerToAddrSpace(op, srcRaw, cbufAddressSpace); FailureOr dst = reinterpretPointerToAddrSpace(op, dstRaw, cbAddressSpace); if (failed(src) || failed(dst)) + { return rewriter.notifyMatchFailure(op, "failed to map cbuf/cb pointer spaces"); + } Type sourceElemType = cast(op.getSource().getType()).getElementType(); unsigned elemBitWidth = pto::getPTOStorageElemBitWidth(sourceElemType); @@ -6363,7 +7360,9 @@ class LowerCopyMatrixCcToGmOpPattern final Value xm = adaptor.getXm(); Value xt = adaptor.getXt(); if (!sourceRaw || !destinationRaw || !xm || !xt) + { return rewriter.notifyMatchFailure(op, "expected converted operands"); + } if (!isa(sourceRaw.getType()) || !isa(destinationRaw.getType())) { @@ -6372,7 +7371,9 @@ class LowerCopyMatrixCcToGmOpPattern final Type i64Ty = rewriter.getI64Type(); if (xm.getType() != i64Ty || xt.getType() != i64Ty) + { return rewriter.notifyMatchFailure(op, "expected i64 xm/xt operands"); + } constexpr unsigned gmAddressSpace = static_cast(pto::AddressSpace::GM); @@ -6382,7 +7383,9 @@ class LowerCopyMatrixCcToGmOpPattern final FailureOr destination = reinterpretPointerToAddrSpace(op, destinationRaw, gmAddressSpace); if (failed(source) || failed(destination)) + { return rewriter.notifyMatchFailure(op, "failed to map cc/gm pointer spaces"); + } StringRef calleeName = buildCopyMatrixCcToGmCallee(op.getContext()); auto funcType = rewriter.getFunctionType( @@ -6414,7 +7417,9 @@ class LowerCopyMatrixCcToBufOpPattern final Value sourceRaw = adaptor.getSource(); Value destinationRaw = adaptor.getDestination(); if (!sourceRaw || !destinationRaw) + { return rewriter.notifyMatchFailure(op, "expected converted operands"); + } if (!isa(sourceRaw.getType()) || !isa(destinationRaw.getType())) return rewriter.notifyMatchFailure(op, "expected LLVM pointer src/dst"); @@ -6430,13 +7435,17 @@ class LowerCopyMatrixCcToBufOpPattern final FailureOr destination = reinterpretPointerToAddrSpace(op, destinationRaw, targetAddressSpace); if (failed(source) || failed(destination)) + { return rewriter.notifyMatchFailure(op, "failed to map cc->buf pointer spaces"); + } Type i64Ty = rewriter.getI64Type(); Value config0 = castIntegerLikeTo(op, adaptor.getConfig0(), i64Ty); Value config1 = castIntegerLikeTo(op, adaptor.getConfig1(), i64Ty); if (!config0 || !config1) + { return rewriter.notifyMatchFailure(op, "failed to cast config operands to i64"); + } FailureOr calleeName = std::is_same_v @@ -6570,7 +7579,9 @@ class LowerHistogramOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { StringRef calleeName = getHistogramCallee(op.getContext()); if (calleeName.empty()) + { return rewriter.notifyMatchFailure(op, "unsupported histogram op"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); @@ -6578,7 +7589,9 @@ class LowerHistogramOpPattern final : public OpConversionPattern { this->getTypeConverter()->convertType(op.getSource().getType()); Type maskType = this->getTypeConverter()->convertType(op.getMask().getType()); if (!resultType || !sourceType || !maskType) + { return rewriter.notifyMatchFailure(op, "failed to convert histogram types"); + } Value acc = adaptor.getAcc(); Value source = adaptor.getSource(); @@ -6715,12 +7728,16 @@ class LowerVselOpPattern final : public OpConversionPattern { FailureOr calleeName = buildVselCallee(op.getContext(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vsel VPTO signature"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); Type maskType = this->getTypeConverter()->convertType(op.getMask().getType()); if (!resultType || !maskType) + { return rewriter.notifyMatchFailure(op, "failed to convert vsel result type"); + } Value src0 = adaptor.getSrc0(); Value src1 = adaptor.getSrc1(); @@ -6755,12 +7772,16 @@ class LowerVdupOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { FailureOr calleeName = buildVdupCallee(op.getContext(), op); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vdup VPTO signature"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); Type maskType = this->getTypeConverter()->convertType(op.getMask().getType()); if (!resultType || !maskType) + { return rewriter.notifyMatchFailure(op, "failed to convert vdup result type"); + } Value mask = adaptor.getMask(); if (!mask || mask.getType() != maskType) @@ -6824,11 +7845,15 @@ class LowerVbrOpPattern final : public OpConversionPattern { buildVbrCallee(op.getContext(), cast(op.getResult().getType()).getElementType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vbr VPTO signature"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert vbr result type"); + } Value scalar = adaptor.getValue(); Type expectedScalarType = @@ -6867,7 +7892,9 @@ class LowerVselrOpPattern final : public OpConversionPattern { FailureOr calleeName = buildVselrCallee(op.getContext(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vselr VPTO signature"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) @@ -6915,7 +7942,9 @@ class LowerVselrOpPattern final : public OpConversionPattern { Value result = call.getResult(0); if (intrinsicResultType != resultType) + { result = rewriter.create(op.getLoc(), resultType, result); + } rewriter.replaceOp(op, ValueRange{result}); return success(); } @@ -7081,7 +8110,9 @@ class LowerUnpackOpPattern final : public OpConversionPattern { Value part = castIntegerLikeTo(op, adaptor.getPart(), rewriter.getI32Type()); if (!part) + { return rewriter.notifyMatchFailure(op, "failed to materialize unpack part"); + } auto funcType = rewriter.getFunctionType(TypeRange{srcType, part.getType()}, TypeRange{resultType}); @@ -7109,17 +8140,23 @@ class LowerVpackOpPattern final : public OpConversionPattern { buildVpackCallee(op.getContext(), op.getSrc().getType(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vpack VPTO signature"); + } Type srcType = this->getTypeConverter()->convertType(op.getSrc().getType()); Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!srcType || !resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert vpack types"); + } auto partImm = parseHiLoPartImmediate(op.getPart()); if (!partImm) + { return rewriter.notifyMatchFailure(op, "unsupported vpack part immediate"); + } Value src = adaptor.getSrc(); if (!src || src.getType() != srcType) { @@ -7238,9 +8275,13 @@ class LowerCmpOpPattern final : public OpConversionPattern { constexpr bool isScalarCompare = std::is_same_v; Type inputType = Type(); if constexpr (isScalarCompare) + { inputType = op.getSrc().getType(); + } else + { inputType = op.getSrc0().getType(); + } FailureOr calleeName = buildVcmpCallee(op.getContext(), inputType, op.getCmpMode(), isScalarCompare); @@ -7302,11 +8343,15 @@ class LowerPltOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Value laneCount = castIntegerLikeTo(op, adaptor.getScalar(), rewriter.getI32Type()); if (!laneCount) + { return rewriter.notifyMatchFailure(op, "failed to materialize plt lane count"); + } SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), resultTypes))) + { return rewriter.notifyMatchFailure(op, "failed to convert plt result types"); + } StringRef calleeName = buildPltCallee(op.getContext()); auto funcType = rewriter.getFunctionType(TypeRange{rewriter.getI32Type()}, @@ -7371,11 +8416,15 @@ class LowerPsetOpPattern final : public OpConversionPattern { (void)adaptor; auto pattern = parsePredicatePatternImmediate(op.getPattern()); if (!pattern) + { return rewriter.notifyMatchFailure(op, "unsupported pset pattern"); + } SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), resultTypes))) + { return rewriter.notifyMatchFailure(op, "failed to convert pset result types"); + } if (isMaskOnlyUsedByOnePointStores(op.getResult())) { auto undef = rewriter.create(op.getLoc(), resultTypes.front()); @@ -7412,11 +8461,15 @@ class LowerPgeOpPattern final : public OpConversionPattern { (void)adaptor; auto pattern = parsePredicatePatternImmediate(op.getPattern()); if (!pattern) + { return rewriter.notifyMatchFailure(op, "unsupported pge pattern"); + } SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), resultTypes))) + { return rewriter.notifyMatchFailure(op, "failed to convert pge result types"); + } if (isMaskOnlyUsedByOnePointStores(op.getResult())) { auto undef = rewriter.create(op.getLoc(), resultTypes.front()); @@ -7455,17 +8508,23 @@ class LowerVldsOpPattern final : public OpConversionPattern { Type ptoResultType = op.getResult().getType(); Type elementType = getElementTypeFromVectorLike(ptoResultType); if (!elementType) + { return rewriter.notifyMatchFailure(op, "unsupported vlds element type"); + } auto offsetBytes = convertElementOffsetToBytes(op, adaptor.getOffset(), elementType); auto basePtr = dyn_cast(adaptor.getSource().getType()); auto dist = parseLoadDistImmediate(op.getDist().value_or("NORM"), elementType); if (failed(offsetBytes) || !basePtr || !dist) + { return rewriter.notifyMatchFailure(op, "failed to materialize vlds operands"); + } SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), resultTypes))) + { return rewriter.notifyMatchFailure(op, "failed to convert vlds result types"); + } bool usePostIntrinsic = static_cast(op.getUpdatedBase()); if (usePostIntrinsic) { @@ -7481,13 +8540,17 @@ class LowerVldsOpPattern final : public OpConversionPattern { ? buildVldsPostCallee(op.getContext(), ptoResultType) : buildVldsCallee(op.getContext(), ptoResultType); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vlds signature"); + } Type callValueType = getPayloadABIType( ptoResultType, resultTypes[0], rewriter.getContext()); SmallVector callResultTypes{callValueType}; if (usePostIntrinsic) + { callResultTypes.push_back(resultTypes[1]); + } Value distValue = getI32Constant(rewriter, op.getLoc(), *dist); Value postValue = getI32Constant(rewriter, op.getLoc(), usePostIntrinsic ? 1 : 0); @@ -7525,7 +8588,9 @@ class LowerVldsx2OpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Type elementType = getElementTypeFromVectorLike(op.getLow().getType()); if (!elementType) + { return rewriter.notifyMatchFailure(op, "unsupported vldsx2 element type"); + } auto offsetBytes = convertElementOffsetToBytes(op, adaptor.getOffset(), elementType); @@ -7549,7 +8614,9 @@ class LowerVldsx2OpPattern final : public OpConversionPattern { buildVldsx2Callee(op.getContext(), op.getLow().getType(), usePostIntrinsic); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vldsx2 signature"); + } Type lowCallType = getPayloadABIType( op.getLow().getType(), resultTypes[0], rewriter.getContext()); @@ -7557,7 +8624,9 @@ class LowerVldsx2OpPattern final : public OpConversionPattern { op.getHigh().getType(), resultTypes[1], rewriter.getContext()); SmallVector callResultTypes{lowCallType, highCallType}; if (usePostIntrinsic) + { callResultTypes.push_back(resultTypes[2]); + } Value distValue = getI32Constant(rewriter, op.getLoc(), *dist); Value postValue = @@ -7601,7 +8670,9 @@ class LowerVsldbOpPattern final : public OpConversionPattern { Value packedStride = packBlockRepeatStride(op, adaptor.getBlockStride(), adaptor.getRepeatStride()); if (!basePtr || !packedStride) + { return rewriter.notifyMatchFailure(op, "failed to materialize vsldb operands"); + } bool usePostIntrinsic = op.getUpdatedBase() != nullptr; SmallVector resultTypes; @@ -7614,13 +8685,17 @@ class LowerVsldbOpPattern final : public OpConversionPattern { op.getResult().getType(), resultTypes[0], rewriter.getContext()); SmallVector callResultTypes{callResultType}; if (usePostIntrinsic) + { callResultTypes.push_back(resultTypes[1]); + } FailureOr calleeName = buildVsldbCallee(op.getContext(), op.getResult().getType(), usePostIntrinsic); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vsldb signature"); + } Value postValue = getI32Constant(rewriter, op.getLoc(), usePostIntrinsic ? 1 : 0); SmallVector args{adaptor.getSource(), packedStride, postValue, @@ -7660,7 +8735,9 @@ class LowerInitAlignOpPattern final (void)adaptor; Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert init_align result type"); + } StringRef calleeName = buildInitAlignCallee(op.getContext()); auto funcType = rewriter.getFunctionType(TypeRange{}, TypeRange{resultType}); @@ -7732,7 +8809,9 @@ class LowerVldusOpPattern final : public OpConversionPattern { ? buildVldusPostCallee(op.getContext(), op.getResult().getType()) : buildVldusCallee(op.getContext(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vldus signature"); + } Type callValueType = getPayloadABIType( op.getResult().getType(), resultTypes[0], rewriter.getContext()); @@ -7752,7 +8831,9 @@ class LowerVldusOpPattern final : public OpConversionPattern { } SmallVector argTypes; for (Value arg : args) + { argTypes.push_back(arg.getType()); + } auto funcType = rewriter.getFunctionType(argTypes, intrinsicResultTypes); auto call = rewriter.create( op.getLoc(), *calleeName, intrinsicResultTypes, args); @@ -7762,7 +8843,9 @@ class LowerVldusOpPattern final : public OpConversionPattern { resultTypes[0], rewriter); SmallVector replacements{loaded, call.getResult(1)}; if (usePostIntrinsic) + { replacements.push_back(call.getResult(2)); + } rewriter.replaceOp(op, replacements); return success(); } @@ -7783,7 +8866,9 @@ class LowerSprclrOpPattern final : public OpConversionPattern { (void)adaptor; auto spr = parseSprImmediate(op.getSpr()); if (!spr) + { return rewriter.notifyMatchFailure(op, "unsupported sprclr target"); + } StringRef calleeName = buildSprclrCallee(op.getContext()); Value sprValue = rewriter.create( @@ -7811,7 +8896,9 @@ class LowerSprStoreOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto spr = parseSprImmediate(op.getSpr()); if (!spr) + { return rewriter.notifyMatchFailure(op, "unsupported spr store target"); + } auto destType = dyn_cast(adaptor.getDestination().getType()); if (!destType || !adaptor.getOffset().getType().isInteger(32)) @@ -7842,9 +8929,13 @@ class LowerSprStoreOpPattern final : public OpConversionPattern { resultTypes, args); state.plannedDecls.push_back(PlannedDecl{calleeName.str(), funcType}); if (usePostIntrinsic) + { rewriter.replaceOp(op, call.getResults()); + } else + { rewriter.eraseOp(op); + } return success(); } @@ -7863,10 +8954,14 @@ class LowerVstsOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Type elementType = getElementTypeFromVectorLike(op.getValue().getType()); if (!elementType) + { return rewriter.notifyMatchFailure(op, "unsupported vsts element type"); + } Type offsetElementType = elementType; if (auto ptrType = dyn_cast(op.getDestination().getType())) + { offsetElementType = ptrType.getElementType(); + } else if (auto memrefType = dyn_cast(op.getDestination().getType())) offsetElementType = memrefType.getElementType(); auto offsetBytes = @@ -7875,14 +8970,18 @@ class LowerVstsOpPattern final : public OpConversionPattern { auto dist = parseStoreDistImmediate(op.getDist().value_or(""), elementType); if (failed(offsetBytes) || !basePtr || !dist) + { return rewriter.notifyMatchFailure(op, "failed to materialize vsts operands"); + } FailureOr calleeName = op.getUpdatedBase() ? buildVstsPostCallee(op.getContext(), op.getValue().getType()) : buildVstsCallee(op.getContext(), op.getValue().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vsts signature"); + } SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), @@ -7911,7 +9010,9 @@ class LowerVstsOpPattern final : public OpConversionPattern { // this dead operand; an LLVM undef is sufficient at this boundary. StringRef distToken = op.getDist().value_or(""); if (isOnePointStoreDist(distToken)) + { mask = rewriter.create(op.getLoc(), mask.getType()); + } SmallVector args{value, adaptor.getDestination(), *offsetBytes, distValue, zero, mask}; auto funcType = rewriter.getFunctionType( @@ -7923,9 +9024,13 @@ class LowerVstsOpPattern final : public OpConversionPattern { resultTypes, args); state.plannedDecls.push_back(PlannedDecl{calleeName->str(), funcType}); if (usePostIntrinsic) + { rewriter.replaceOp(op, call.getResults()); + } else + { rewriter.eraseOp(op); + } return success(); } @@ -7947,7 +9052,9 @@ class LowerVsstbOpPattern final : public OpConversionPattern { Value packedStride = packBlockRepeatStride(op, adaptor.getBlockStride(), adaptor.getRepeatStride()); if (!basePtr || !packedStride) + { return rewriter.notifyMatchFailure(op, "failed to materialize vsstb operands"); + } SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), @@ -7967,7 +9074,9 @@ class LowerVsstbOpPattern final : public OpConversionPattern { FailureOr calleeName = buildVsstbCallee( op.getContext(), op.getValue().getType(), usePostIntrinsic); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vsstb signature"); + } Value zeroValue = getI32Constant(rewriter, op.getLoc(), usePostIntrinsic ? 1 : 0); Value value = castToPayloadABI( op.getLoc(), adaptor.getValue(), op.getValue().getType(), rewriter); @@ -7982,9 +9091,13 @@ class LowerVsstbOpPattern final : public OpConversionPattern { resultTypes, args); state.plannedDecls.push_back(PlannedDecl{calleeName->str(), funcType}); if (usePostIntrinsic) + { rewriter.replaceOp(op, call.getResults()); + } else + { rewriter.eraseOp(op); + } return success(); } @@ -8004,7 +9117,9 @@ class LowerVstsx2OpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Type elementType = getElementTypeFromVectorLike(op.getLow().getType()); if (!elementType) + { return rewriter.notifyMatchFailure(op, "unsupported vstsx2 element type"); + } auto offsetBytes = convertElementOffsetToBytes(op, adaptor.getOffset(), elementType); @@ -8019,7 +9134,9 @@ class LowerVstsx2OpPattern final : public OpConversionPattern { FailureOr calleeName = buildVstsx2Callee(op.getContext(), op.getLow().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vstsx2 signature"); + } Value distValue = getI32Constant(rewriter, op.getLoc(), *dist); Value zeroValue = getI32Constant(rewriter, op.getLoc(), 0); @@ -8056,13 +9173,19 @@ class LowerPstuOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { FailureOr calleeName = buildPstuCallee(op.getContext(), op); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported pstu signature"); + } SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), resultTypes))) + { return rewriter.notifyMatchFailure(op, "failed to convert pstu result types"); + } if (resultTypes.size() != 2) + { return rewriter.notifyMatchFailure(op, "unexpected converted pstu result arity"); + } auto baseType = dyn_cast(adaptor.getBase().getType()); if (!baseType || adaptor.getAlignIn().getType() != resultTypes[0] || @@ -8098,11 +9221,15 @@ class LowerVstusOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Type elementType = getElementTypeFromVectorLike(op.getValue().getType()); if (!elementType) + { return rewriter.notifyMatchFailure(op, "unsupported vstus element type"); + } auto offsetBytes = convertElementOffsetToBytes(op, adaptor.getOffset(), elementType); if (failed(offsetBytes)) + { return rewriter.notifyMatchFailure(op, "failed to convert vstus offset"); + } SmallVector resultTypes; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), @@ -8124,7 +9251,9 @@ class LowerVstusOpPattern final : public OpConversionPattern { calleeName = buildVstusPostCallee(op.getContext(), op.getValue().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vstus signature"); + } Value value = castToPayloadABI( op.getLoc(), adaptor.getValue(), op.getValue().getType(), rewriter); SmallVector args{value, adaptor.getBase(), *offsetBytes, @@ -8155,7 +9284,9 @@ class LowerVsturOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto postMode = parsePostModeImmediate(op.getMode()); if (!postMode) + { return rewriter.notifyMatchFailure(op, "unsupported vstur mode immediate"); + } Type resultType = this->getTypeConverter()->convertType(op.getAlignOut().getType()); auto baseType = dyn_cast(adaptor.getBase().getType()); @@ -8240,7 +9371,9 @@ class LowerVstasOpPattern final : public OpConversionPattern { auto offsetBytes = convertElementOffsetToBytes(op, adaptor.getOffset(), dstType.getElementType()); if (failed(offsetBytes)) + { return rewriter.notifyMatchFailure(op, "failed to convert vstas offset"); + } bool usePostIntrinsic = op.getUpdatedBase() != nullptr; SmallVector resultTypes; @@ -8264,9 +9397,13 @@ class LowerVstasOpPattern final : public OpConversionPattern { resultTypes, args); state.plannedDecls.push_back(PlannedDecl{calleeName.str(), funcType}); if (usePostIntrinsic) + { rewriter.replaceOp(op, call.getResults()); + } else + { rewriter.eraseOp(op); + } return success(); } @@ -8293,20 +9430,26 @@ class LowerVgather2OpPattern final Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert vgather2 result type"); + } FailureOr calleeName = buildVgather2Callee(op.getContext(), op.getSource().getType(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vgather2 signature"); + } Value offsets = adaptor.getOffsets(); FailureOr offsetsCarrierType = getVgather2OffsetsCarrierType( rewriter, op.getSource().getType(), op.getResult().getType(), offsets.getType()); if (failed(offsetsCarrierType)) + { return rewriter.notifyMatchFailure(op, "unsupported vgather2 offsets carrier"); + } if (offsets.getType() != *offsetsCarrierType) offsets = rewriter.create(op.getLoc(), *offsetsCarrierType, offsets); @@ -8347,7 +9490,9 @@ class LowerVgather2BcOpPattern final FailureOr calleeName = buildVgather2BcCallee(op.getContext(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vgather2_bc signature"); + } auto funcType = rewriter.getFunctionType( TypeRange{adaptor.getSource().getType(), adaptor.getOffsets().getType(), @@ -8385,7 +9530,9 @@ class LowerVgatherbOpPattern final FailureOr calleeName = buildVgatherbCallee(op.getContext(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vgatherb signature"); + } auto funcType = rewriter.getFunctionType( TypeRange{adaptor.getSource().getType(), adaptor.getOffsets().getType(), @@ -8424,12 +9571,16 @@ class LowerVscatterOpPattern final FailureOr calleeName = buildVscatterCallee(op.getContext(), op.getValue().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vscatter signature"); + } FailureOr offsetsCarrierType = getVscatterOffsetsCarrierType( adaptor.getOffsets().getType()); if (failed(offsetsCarrierType)) + { return rewriter.notifyMatchFailure(op, "unsupported vscatter offsets carrier"); + } auto funcType = rewriter.getFunctionType( TypeRange{adaptor.getValue().getType(), adaptor.getDestination().getType(), @@ -8459,16 +9610,22 @@ class LowerVaxpyOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Type elemType = getElementTypeFromVectorLike(op.getResult().getType()); if (!elemType) + { return rewriter.notifyMatchFailure(op, "unsupported vaxpy signature"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert vaxpy result type"); + } FailureOr calleeName = buildVaxpyCallee(op.getContext(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vaxpy callee"); + } auto funcType = rewriter.getFunctionType( TypeRange{adaptor.getSrc1().getType(), adaptor.getSrc0().getType(), @@ -8500,14 +9657,18 @@ class LowerVmulscvtOpPattern final ConversionPatternRewriter &rewriter) const override { auto roundMode = parseRoundModeImmediate(op.getRnd()); if (!roundMode) + { return rewriter.notifyMatchFailure(op, "vmulscvt requires valid rnd attr"); + } if (*roundMode != 1) return rewriter.notifyMatchFailure( op, "current vmulscvt lowering only supports rnd A"); auto part = parsePartImmediate(op.getPart()); if (!part) + { return rewriter.notifyMatchFailure(op, "unsupported vmulscvt part"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); @@ -8519,7 +9680,9 @@ class LowerVmulscvtOpPattern final buildVmulscvtCallee(op.getContext(), op.getInput().getType(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vmulscvt signature"); + } Value partValue = getI32Constant(rewriter, op.getLoc(), *part); auto funcType = rewriter.getFunctionType( @@ -8550,16 +9713,22 @@ class LowerVciOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto order = parseOrderImmediate(op.getOrder().value_or("ASC")); if (!order) + { return rewriter.notifyMatchFailure(op, "unsupported vci order"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert vci result type"); + } FailureOr calleeName = buildVciCallee(op.getContext(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vci callee"); + } Value indexValue = adaptor.getIndex(); Type resultElemType = @@ -8607,17 +9776,23 @@ class LowerVexpdifOpPattern final ConversionPatternRewriter &rewriter) const override { auto part = parsePartImmediate(op.getPart()); if (!part) + { return rewriter.notifyMatchFailure(op, "unsupported vexpdif signature"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert vexpdif result type"); + } FailureOr calleeName = buildVexpdifCallee(op.getContext(), op.getInput().getType(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vexpdif callee"); + } Value partValue = getI32Constant(rewriter, op.getLoc(), *part); auto funcType = rewriter.getFunctionType( @@ -8659,11 +9834,15 @@ class LowerVbitsortOpPattern final FailureOr config = packVbitsortConfig(op, adaptor.getRepeatTimes()); if (failed(config)) + { return rewriter.notifyMatchFailure(op, "failed to pack vbitsort config"); + } FailureOr calleeName = buildVbitsortCallee(op.getContext(), op); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vbitsort signature"); + } auto funcType = rewriter.getFunctionType( TypeRange{adaptor.getDestination().getType(), adaptor.getSource().getType(), @@ -8718,11 +9897,15 @@ class LowerVmrgsort4OpPattern final FailureOr dst = reinterpretPointerToAddrSpace(op, adaptor.getDestination(), 6); if (failed(dst)) + { return rewriter.notifyMatchFailure(op, "failed to normalize vmrgsort4 destination"); + } FailureOr calleeName = buildVmrgsort4Callee(op.getContext(), op); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vmrgsort4 signature"); + } auto funcType = rewriter.getFunctionType( TypeRange{(*dst).getType(), (*packedSrc).getType(), @@ -8751,11 +9934,15 @@ class LowerVcvtOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { FailureOr contract = buildVcvtContract(op); if (failed(contract)) + { return rewriter.notifyMatchFailure(op, "unsupported vcvt type pair"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert vcvt result type"); + } SmallVector callArgs; SmallVector argTypes; @@ -8768,7 +9955,9 @@ class LowerVcvtOpPattern final : public OpConversionPattern { auto roundMode = op.getRndAttr() ? parseRoundModeImmediate(*op.getRnd()) : std::nullopt; if (!roundMode) + { return rewriter.notifyMatchFailure(op, "vcvt requires valid rnd attr"); + } Value roundValue = getI32Constant(rewriter, op.getLoc(), *roundMode); callArgs.push_back(roundValue); argTypes.push_back(roundValue.getType()); @@ -8779,7 +9968,9 @@ class LowerVcvtOpPattern final : public OpConversionPattern { auto saturation = op.getSatAttr() ? parseSaturationImmediate(*op.getSat()) : std::nullopt; if (!saturation) + { return rewriter.notifyMatchFailure(op, "vcvt requires valid sat attr"); + } Value satValue = getI32Constant(rewriter, op.getLoc(), *saturation); callArgs.push_back(satValue); argTypes.push_back(satValue.getType()); @@ -8788,21 +9979,31 @@ class LowerVcvtOpPattern final : public OpConversionPattern { if ((*contract).satBeforeRnd) { if ((*contract).requiresSat && failed(appendSatArg())) + { return failure(); + } if ((*contract).requiresRnd && failed(appendRndArg())) + { return failure(); + } } else { if ((*contract).requiresRnd && failed(appendRndArg())) + { return failure(); + } if ((*contract).requiresSat && failed(appendSatArg())) + { return failure(); + } } if ((*contract).requiresPart) { auto part = op.getPartAttr() ? parseVcvtPartImmediate(*op.getPart()) : std::nullopt; if (!part) + { return rewriter.notifyMatchFailure(op, "vcvt requires valid part attr"); + } Value partValue = getI32Constant(rewriter, op.getLoc(), *part); callArgs.push_back(partValue); argTypes.push_back(partValue.getType()); @@ -8877,16 +10078,22 @@ class LowerVtrcOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto roundMode = parseRoundModeImmediate(op.getRoundMode()); if (!roundMode) + { return rewriter.notifyMatchFailure(op, "unsupported vtrc signature"); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert vtrc result type"); + } FailureOr calleeName = buildVtrcCallee(op.getContext(), op.getResult().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported vtrc callee"); + } Value roundValue = getI32Constant(rewriter, op.getLoc(), *roundMode); auto funcType = rewriter.getFunctionType( @@ -8960,9 +10167,13 @@ class LowerPredicateStoreOpPattern final : public OpConversionPattern { resultTypes, args); state.plannedDecls.push_back(PlannedDecl{calleeName.str(), funcType}); if (usePostIntrinsic) + { rewriter.replaceOp(op, call.getResults()); + } else + { rewriter.eraseOp(op); + } return success(); } @@ -9140,12 +10351,16 @@ class LowerStoreVfSimtInfoOpPattern final Value dimY = adaptor.getDimY(); Value dimX = adaptor.getDimX(); if (!dimZ || !dimY || !dimX) + { return rewriter.notifyMatchFailure(op, "missing converted SIMT dims"); + } auto i64Type = rewriter.getI64Type(); auto castToI64 = [&](Value value) -> Value { if (value.getType().isInteger(64)) + { return value; + } return rewriter.create(loc, i64Type, value).getResult(); }; @@ -9236,7 +10451,9 @@ static std::string buildSimtKeepResumeConstraints( llvm::raw_string_ostream os(result); for (auto [index, physicalReg] : llvm::enumerate(physicalRegs)) { if (index != 0) + { os << ","; + } if (physicalReg.registerCount == 2) os << "={TPERL" << physicalReg.baseRegister / 2 << "}"; else @@ -9244,7 +10461,9 @@ static std::string buildSimtKeepResumeConstraints( } if (tieInputs) { for (size_t index = 0; index < physicalRegs.size(); ++index) + { os << "," << index; + } } return os.str(); } @@ -9255,7 +10474,9 @@ static SmallVector collectConsecutiveOps(OpT first) { for (Operation *cur = first.getOperation(); cur; cur = cur->getNextNode()) { auto typed = dyn_cast(cur); if (!typed) + { break; + } ops.push_back(typed); } return ops; @@ -9269,13 +10490,19 @@ static bool hasPreviousSameOp(Operation *op) { static std::optional getSimtKeepResumeBitWidth(Type type) { if (auto intType = dyn_cast(type)) { if (intType.getWidth() <= 64) + { return intType.getWidth(); + } return std::nullopt; } if (type.isF16() || type.isBF16()) + { return 16; + } if (type.isF32()) + { return 32; + } return std::nullopt; } @@ -9289,13 +10516,19 @@ static Value packSimtKeepResumePayload(Location loc, Value value, Type intType = rewriter.getIntegerType(*width); Value bits = value; if (!isa(type)) + { bits = rewriter.create(loc, intType, value); + } else if (bits.getType() != intType) bits = rewriter.create(loc, intType, bits); if (*width < 32) + { return rewriter.create(loc, rewriter.getI32Type(), bits); + } if (*width == 32 && bits.getType() != rewriter.getI32Type()) + { return rewriter.create(loc, rewriter.getI32Type(), bits); + } return bits; } @@ -9309,13 +10542,17 @@ static Value unpackSimtKeepResumePayload(Location loc, Value value, Type intType = rewriter.getIntegerType(*width); Value bits = value; if (*width < 32) + { bits = rewriter.create(loc, intType, bits); + } else if (bits.getType() != intType) bits = rewriter.create(loc, intType, bits); if (isa(resultType)) { if (bits.getType() == resultType) + { return bits; + } return rewriter.create(loc, resultType, bits); } return rewriter.create(loc, resultType, bits); @@ -9333,15 +10570,21 @@ computeSimtKeepResumePhysicalRegs( physicalRegs.reserve(logicalSlots.size()); for (auto [slot, registerCount] : logicalSlots) { if (slot < 0 || slot >= 123) + { return failure(); + } if (registerCount == 2 && ((slot % 2) != 0 || slot + 1 >= 123)) + { return failure(); + } // Slots are user-assigned storage words, not dense ordinals in the current // keep/resume group. This keeps a consumer that resumes only a subset of // slots from changing where the remaining slots are read from. int64_t baseRegister = 4 + slot; if (baseRegister + static_cast(registerCount) - 1 > 126) + { return failure(); + } physicalRegs.push_back({baseRegister, registerCount}); } return physicalRegs; @@ -9349,9 +10592,13 @@ computeSimtKeepResumePhysicalRegs( static bool isValidSimtKeepResumeSlot(int64_t slot, unsigned registerCount) { if (slot < 0 || slot >= 123) + { return false; + } if (registerCount == 2 && ((slot % 2) != 0 || slot + 1 >= 123)) + { return false; + } return true; } @@ -9376,7 +10623,9 @@ class LowerKeepOpPattern final : public OpConversionPattern { for (pto::KeepOp keep : keepOps) { Value payload = rewriter.getRemappedValue(keep.getPayload()); if (!payload) + { return rewriter.notifyMatchFailure(keep, "payload is not remapped"); + } payload = packSimtKeepResumePayload(keep.getLoc(), payload, rewriter); if (!payload) return rewriter.notifyMatchFailure( @@ -9410,7 +10659,9 @@ class LowerKeepOpPattern final : public OpConversionPattern { LLVM::AsmDialectAttr::get(op.getContext(), LLVM::AsmDialect::AD_ATT), ArrayAttr{}); for (pto::KeepOp keep : llvm::reverse(keepOps)) + { rewriter.eraseOp(keep); + } return success(); } }; @@ -9471,7 +10722,9 @@ class LowerResumeOpPattern final : public OpConversionPattern { Value result = unpackSimtKeepResumePayload(op.getLoc(), asmOp.getRes(), resultType, rewriter); if (!result) + { return rewriter.notifyMatchFailure(op, "failed to unpack result"); + } rewriter.replaceOp(op, result); return success(); } @@ -9486,11 +10739,15 @@ class LowerResumeOpPattern final : public OpConversionPattern { Value result = unpackSimtKeepResumePayload( resume.getLoc(), extract.getRes(), resultType, rewriter); if (!result) + { return rewriter.notifyMatchFailure(resume, "failed to unpack result"); + } results.push_back(result); } for (auto [resume, result] : llvm::zip(resumeOps, results)) + { rewriter.replaceOp(resume, result); + } return success(); } }; @@ -9536,7 +10793,9 @@ class LowerPipeEventSyncOpPattern final : public OpConversionPattern { auto dst = parsePipeImmediate(stringifyPIPE(op.getDstPipe().getPipe())); auto event = parseEventImmediate(stringifyEVENT(op.getEventId().getEvent())); if (!src || !dst || !event) + { return rewriter.notifyMatchFailure(op, "unsupported sync immediate"); + } StringRef calleeName = buildSyncCallee(op.getContext()); Value srcValue = getI64Constant(rewriter, op.getLoc(), *src); @@ -9571,7 +10830,9 @@ class LowerPipeEventDynSyncOpPattern final : public OpConversionPattern auto src = parsePipeImmediate(stringifyPIPE(op.getSrcPipe().getPipe())); auto dst = parsePipeImmediate(stringifyPIPE(op.getDstPipe().getPipe())); if (!src || !dst) + { return rewriter.notifyMatchFailure(op, "unsupported sync pipe"); + } StringRef calleeName = buildSyncCallee(op.getContext()); Value srcValue = getI64Constant(rewriter, op.getLoc(), *src); @@ -9579,14 +10840,18 @@ class LowerPipeEventDynSyncOpPattern final : public OpConversionPattern Value eventIdValue = adaptor.getEventId(); if (!eventIdValue) + { return rewriter.notifyMatchFailure(op, "missing event_id operand"); + } Value eventValue = eventIdValue; while (eventValue.getDefiningOp()) { auto unrealizedCast = dyn_cast(eventValue.getDefiningOp()); if (!unrealizedCast || unrealizedCast.getInputs().size() != 1) + { break; + } eventValue = unrealizedCast.getInputs()[0]; } @@ -9632,7 +10897,9 @@ class LowerInterCoreSyncOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto pipe = parsePipeImmediate(stringifyPIPE(op.getPipe().getPipe())); if (!pipe) + { return rewriter.notifyMatchFailure(op, "unsupported inter-core sync pipe"); + } Value pipeValue = getI64Constant(rewriter, op.getLoc(), *pipe); Value eventValue; @@ -9685,7 +10952,9 @@ class LowerBarrierOpPattern final : public OpConversionPattern { auto pipe = parsePipeImmediate(stringifyPIPE(op.getPipe().getPipe())); if (!pipe) + { return rewriter.notifyMatchFailure(op, "unsupported barrier pipe"); + } StringRef calleeName = buildSyncCallee(op.getContext()); Value pipeValue = getI64Constant(rewriter, op.getLoc(), *pipe); @@ -9789,7 +11058,9 @@ class LowerDcciOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto ptrType = dyn_cast(adaptor.getPtr().getType()); if (!ptrType) + { return rewriter.notifyMatchFailure(op, "expected LLVM pointer operand"); + } bool hasDst = static_cast(op.getDstAttr()); StringRef calleeName = @@ -9845,7 +11116,9 @@ class LowerBufSyncOpPattern final : public OpConversionPattern { auto pipeImm = parsePipeImmediate(stringifyPIPE(pipe)); if (!pipeImm) + { return rewriter.notifyMatchFailure(op, "unsupported buffer sync pipe"); + } StringRef calleeName = buildSyncCallee(op.getContext()); Value pipeValue = getI64Constant(rewriter, op.getLoc(), *pipeImm); @@ -9896,7 +11169,9 @@ class LowerBufDynSyncOpPattern final auto pipeImm = parsePipeImmediate(stringifyPIPE(pipe)); if (!pipeImm) + { return rewriter.notifyMatchFailure(op, "unsupported buffer sync pipe"); + } Value pipeValue = getI64Constant(rewriter, op.getLoc(), *pipeImm); Value bufIdDyn = adaptor.getBufId(); @@ -9997,7 +11272,9 @@ class LowerBlockRuntimeQueryOpPattern final Value result = call.getResult(0); if (isSimtEntry) + { result = rewriter.create(op.getLoc(), resultType, result); + } rewriter.replaceOp(op, result); return success(); } @@ -10018,11 +11295,15 @@ class LowerVoteOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert vote result type"); + } Type predType = this->getTypeConverter()->convertType(op.getPred().getType()); if (!predType || predType != rewriter.getI1Type()) + { return rewriter.notifyMatchFailure(op, "failed to convert vote predicate type"); + } StringRef calleeName = buildVoteCallee(op.getContext()); auto funcType = rewriter.getFunctionType(TypeRange{predType}, TypeRange{resultType}); @@ -10050,16 +11331,22 @@ class LowerShuffleOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert shuffle result type"); + } Type valueType = this->getTypeConverter()->convertType(op.getValue().getType()); if (!valueType || valueType != resultType) + { return rewriter.notifyMatchFailure(op, "unexpected converted shuffle operand type"); + } FailureOr calleeName = buildShuffleCallee(op.getContext(), op.getValue().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported shuffle VPTO signature"); + } IntegerAttr widthAttr = op.getWidthAttr(); Value controlValue; @@ -10078,7 +11365,9 @@ class LowerShuffleOpPattern final : public OpConversionPattern { controlMask = 0x1f; } if (!controlValue) + { return rewriter.notifyMatchFailure(op, "missing shuffle control operand"); + } Value control = buildShuffleControlValue( rewriter, op.getLoc(), controlValue, widthAttr.getInt(), controlMask); @@ -10110,16 +11399,22 @@ class LowerReduxOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert redux result type"); + } Type valueType = this->getTypeConverter()->convertType(op.getValue().getType()); if (!valueType || valueType != resultType) + { return rewriter.notifyMatchFailure(op, "unexpected converted redux operand type"); + } FailureOr calleeName = buildReduxCallee( op.getContext(), op.getValue().getType(), op.getSignednessAttr()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported redux VPTO signature"); + } auto funcType = rewriter.getFunctionType(TypeRange{resultType}, TypeRange{resultType}); @@ -10154,13 +11449,17 @@ class LowerAtomicBinaryOpPattern final : public OpConversionPattern { Type ptrType = this->getTypeConverter()->convertType(op.getPtr().getType()); if (!ptrType) + { return rewriter.notifyMatchFailure(op, "failed to convert atomic pointer type"); + } FailureOr calleeName = buildAtomicCallee( op.getContext(), op.getPtr().getType(), op.getValue().getType(), op.getSignednessAttr()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported atomic VPTO signature"); + } auto funcType = rewriter.getFunctionType( TypeRange{ptrType, valueType, rewriter.getI32Type()}, @@ -10204,13 +11503,17 @@ class LowerAtomicCasOpPattern final Type ptrType = this->getTypeConverter()->convertType(op.getPtr().getType()); if (!ptrType) + { return rewriter.notifyMatchFailure(op, "failed to convert atomic pointer type"); + } FailureOr calleeName = buildAtomicCallee( op.getContext(), op.getPtr().getType(), op.getValue().getType(), op.getSignednessAttr()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported atomic CAS signature"); + } auto funcType = rewriter.getFunctionType( TypeRange{ptrType, compareType, valueType, rewriter.getI32Type()}, @@ -10306,7 +11609,9 @@ class LowerMulhiOpPattern final : public OpConversionPattern { buildMulhiCallee(op.getContext(), op.getResult().getType(), pto::Signedness::Unsigned); if (failed(unsignedCalleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported mul64hi signature"); + } auto funcType = rewriter.getFunctionType(TypeRange{lhsType, rhsType}, TypeRange{resultType}); @@ -10352,7 +11657,9 @@ class LowerMulI32ToI64OpPattern final Type lhsType = getTypeConverter()->convertType(op.getLhs().getType()); Type rhsType = getTypeConverter()->convertType(op.getRhs().getType()); if (!resultType || !lhsType || !rhsType) + { return rewriter.notifyMatchFailure(op, "unexpected mul_i32toi64 type"); + } FailureOr calleeName = buildMulI32ToI64Callee(op.getContext(), @@ -10387,12 +11694,16 @@ class LowerSqrtOpPattern final : public OpConversionPattern { Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); Type valueType = this->getTypeConverter()->convertType(op.getValue().getType()); if (!resultType || !valueType || valueType != resultType) + { return rewriter.notifyMatchFailure(op, "unexpected sqrt operand/result type"); + } FailureOr calleeName = buildSqrtCallee(op.getContext(), op.getValue().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported sqrt VPTO signature"); + } auto funcType = rewriter.getFunctionType(TypeRange{valueType}, TypeRange{resultType}); @@ -10422,12 +11733,16 @@ class LowerUnaryScalarMathOpPattern final : public OpConversionPattern Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); Type valueType = this->getTypeConverter()->convertType(op.getValue().getType()); if (!resultType || !valueType || valueType != resultType) + { return rewriter.notifyMatchFailure(op, "unexpected unary scalar math type"); + } FailureOr calleeName = buildUnaryScalarMathCallee(op.getContext(), op.getValue().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported unary scalar math signature"); + } auto funcType = rewriter.getFunctionType(TypeRange{valueType}, TypeRange{resultType}); @@ -10464,7 +11779,9 @@ class LowerBinaryScalarMathOpPattern final : public OpConversionPattern calleeName = buildBinaryScalarMathCallee(op.getContext(), op.getLhs().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported binary scalar math signature"); + } auto funcType = rewriter.getFunctionType(TypeRange{lhsType, rhsType}, TypeRange{resultType}); @@ -10500,7 +11817,9 @@ class LowerFmaOpPattern final : public OpConversionPattern { FailureOr calleeName = buildFmaCallee(op.getContext(), op.getLhs().getType()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported fma scalar signature"); + } auto funcType = rewriter.getFunctionType(TypeRange{lhsType, rhsType, accType}, TypeRange{resultType}); @@ -10535,7 +11854,9 @@ class LowerConvertOpPattern final : public OpConversionPattern { buildConvertCallee(op.getContext(), op.getSrc().getType(), op.getDst().getType(), op.getSignednessAttr()); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported convert signature"); + } Value rounding = getI32Constant( rewriter, op.getLoc(), static_cast(op.getRounding())); @@ -10619,7 +11940,9 @@ class LowerBinaryI64PureOpPattern final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "failed to convert result type"); + } StringRef calleeName = buildBinaryI64PureCallee(op.getContext()); auto funcType = @@ -10647,7 +11970,9 @@ class ConvertVPTOUnrealizedCastOp final matchAndRewrite(UnrealizedConversionCastOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { if (op->getNumOperands() != 1 || op->getNumResults() != 1) + { return rewriter.notifyMatchFailure(op, "expected single-operand single-result cast"); + } if (!hasVPTOConvertibleType(op->getOperandTypes()) && !hasVPTOConvertibleType(op->getResultTypes())) return rewriter.notifyMatchFailure(op, "no VPTO convertible types"); @@ -10655,11 +11980,15 @@ class ConvertVPTOUnrealizedCastOp final Type convertedResultType = getTypeConverter()->convertType(op.getResult(0).getType()); if (!convertedResultType) + { return rewriter.notifyMatchFailure(op, "could not convert result type"); + } Value input = adaptor.getOperands().front(); if (input.getType() != convertedResultType) + { return rewriter.notifyMatchFailure(op, "input type does not match converted result type"); + } rewriter.replaceOp(op, input); return success(); @@ -10678,7 +12007,9 @@ class ConvertPtoTileBufAddrOp final getTypeConverter()->convertType(op.getResult().getType()); auto llvmPtrType = dyn_cast(convertedResultType); if (!llvmPtrType) + { return rewriter.notifyMatchFailure(op, "expected LLVM pointer result"); + } Value input = adaptor.getSrc(); if (isa(input.getType())) { @@ -10744,12 +12075,16 @@ class ConvertPtoStructGetOp final ConversionPatternRewriter &rewriter) const override { Type resultType = getTypeConverter()->convertType(op.getValue().getType()); if (!resultType) + { return rewriter.notifyMatchFailure(op, "could not convert result type"); + } FailureOr address = getVPTOStructFieldAddress( rewriter, op.getLoc(), adaptor.getS(), cast(op.getS().getType()), op.getPath()); if (failed(address)) + { return rewriter.notifyMatchFailure(op, "invalid struct field path"); + } rewriter.replaceOpWithNewOp( op, resultType, *address, getNaturalByteAlignment(resultType)); return success(); @@ -10768,7 +12103,9 @@ class ConvertPtoStructSetOp final rewriter, op.getLoc(), adaptor.getS(), cast(op.getS().getType()), op.getPath()); if (failed(address)) + { return rewriter.notifyMatchFailure(op, "invalid struct field path"); + } rewriter.replaceOpWithNewOp( op, adaptor.getValue(), *address, getNaturalByteAlignment(adaptor.getValue().getType())); @@ -10795,7 +12132,9 @@ class ConvertArithSelectOp final : public OpConversionPattern { Type convertedResultType = getTypeConverter()->convertType(op.getResult().getType()); if (!convertedResultType) + { return rewriter.notifyMatchFailure(op, "failed to convert result type"); + } Value trueValue = adaptor.getTrueValue(); Value falseValue = adaptor.getFalseValue(); @@ -10821,7 +12160,9 @@ class ConvertPtoAddPtrOp final : public OpConversionPattern { Type convertedResultType = getTypeConverter()->convertType(op.getResult().getType()); auto llvmPtrType = dyn_cast(convertedResultType); if (!llvmPtrType) + { return rewriter.notifyMatchFailure(op, "expected LLVM pointer result type"); + } Value offset = adaptor.getOffset(); if (offset.getType().isIndex()) @@ -10906,7 +12247,9 @@ class ConvertPtoLoadScalarOp final ConversionPatternRewriter &rewriter) const override { auto llvmPtrType = dyn_cast(adaptor.getPtr().getType()); if (!llvmPtrType) + { return rewriter.notifyMatchFailure(op, "expected LLVM pointer operand"); + } Type convertedValueType = getTypeConverter()->convertType(op.getValue().getType()); @@ -10945,7 +12288,9 @@ class ConvertPtoStoreScalarOp final ConversionPatternRewriter &rewriter) const override { auto llvmPtrType = dyn_cast(adaptor.getPtr().getType()); if (!llvmPtrType) + { return rewriter.notifyMatchFailure(op, "expected LLVM pointer operand"); + } Value offset = adaptor.getOffset(); if (offset.getType().isIndex()) @@ -10979,12 +12324,16 @@ class ConvertPtoLoadOp final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto llvmPtrType = dyn_cast(adaptor.getPtr().getType()); if (!llvmPtrType) + { return rewriter.notifyMatchFailure(op, "expected LLVM pointer operand"); + } Type convertedValueType = getTypeConverter()->convertType(op.getValue().getType()); if (!convertedValueType) + { return rewriter.notifyMatchFailure(op, "could not convert load result type"); + } Value offset = adaptor.getOffset(); if (offset.getType().isIndex()) @@ -11012,23 +12361,37 @@ static Type getLdgCallResultType(Type valueType, Type convertedValueType, if (auto intType = dyn_cast(valueType)) { unsigned width = intType.getWidth(); if (width == 8 || width == 16) + { return rewriter.getI32Type(); + } return convertedValueType; } if (valueType.isF16() || valueType.isBF16() || valueType.isF32()) + { return rewriter.getI32Type(); + } if (valueType.isF64()) + { return rewriter.getI64Type(); + } if (pto::isPTOFloat8Type(valueType) || pto::isPTOHiFloat8Type(valueType)) + { return rewriter.getI32Type(); + } if (pto::isPTOPackedLdgStgVectorType(valueType)) { unsigned totalBits = pto::getPTOPackedLdgStgTotalBits(valueType); if (totalBits == 16) + { return rewriter.getI32Type(); + } if (totalBits == 32) + { return rewriter.getI32Type(); + } if (totalBits == 64) + { return rewriter.getI64Type(); + } } return convertedValueType; } @@ -11082,12 +12445,16 @@ class ConvertPtoLdgOp final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto llvmPtrType = dyn_cast(adaptor.getPtr().getType()); if (!llvmPtrType) + { return rewriter.notifyMatchFailure(op, "expected LLVM pointer operand"); + } Type convertedValueType = getTypeConverter()->convertType(op.getValue().getType()); if (!convertedValueType) + { return rewriter.notifyMatchFailure(op, "could not convert ldg result type"); + } Value offset = adaptor.getOffset(); if (offset.getType().isIndex()) @@ -11108,7 +12475,9 @@ class ConvertPtoLdgOp final : public OpConversionPattern { op, elemPtr, static_cast(ptrTy.getMemorySpace().getAddressSpace())); if (failed(ptr)) + { return rewriter.notifyMatchFailure(op, "failed to map ldg pointer"); + } pto::L1Cache l1cache = op.getL1cacheAttr() ? op.getL1cacheAttr().getValue() @@ -11116,7 +12485,9 @@ class ConvertPtoLdgOp final : public OpConversionPattern { FailureOr calleeName = buildL1CacheLoadCallee( op.getContext(), op.getValue().getType(), l1cache); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported ldg signature"); + } pto::LdL2Cache mode = op.getL2cacheAttr() ? op.getL2cacheAttr().getValue() @@ -11154,7 +12525,9 @@ class ConvertPtoStoreOp final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto llvmPtrType = dyn_cast(adaptor.getPtr().getType()); if (!llvmPtrType) + { return rewriter.notifyMatchFailure(op, "expected LLVM pointer operand"); + } Value offset = adaptor.getOffset(); if (offset.getType().isIndex()) @@ -11181,9 +12554,13 @@ static Value convertStgValue(Location loc, Type valueType, Value value, if (auto intType = dyn_cast(valueType)) { unsigned width = intType.getWidth(); if (width == 8) + { return rewriter.create(loc, rewriter.getI32Type(), value); + } if (width == 16) + { return rewriter.create(loc, rewriter.getF16Type(), value); + } return value; } @@ -11193,11 +12570,17 @@ static Value convertStgValue(Location loc, Type valueType, Value value, return rewriter.create(loc, rewriter.getI32Type(), payload); } if (valueType.isBF16()) + { return rewriter.create(loc, rewriter.getF16Type(), value); + } if (valueType.isF32()) + { return rewriter.create(loc, rewriter.getI32Type(), value); + } if (valueType.isF64()) + { return rewriter.create(loc, rewriter.getI64Type(), value); + } if (pto::isPTOPackedLdgStgVectorType(valueType)) { unsigned totalBits = pto::getPTOPackedLdgStgTotalBits(valueType); if (totalBits == 16) @@ -11225,7 +12608,9 @@ class ConvertPtoStgOp final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto llvmPtrType = dyn_cast(adaptor.getPtr().getType()); if (!llvmPtrType) + { return rewriter.notifyMatchFailure(op, "expected LLVM pointer operand"); + } Value offset = adaptor.getOffset(); if (offset.getType().isIndex()) @@ -11246,7 +12631,9 @@ class ConvertPtoStgOp final : public OpConversionPattern { op, elemPtr, static_cast(ptrTy.getMemorySpace().getAddressSpace())); if (failed(ptr)) + { return rewriter.notifyMatchFailure(op, "failed to map stg pointer"); + } pto::L1Cache l1cache = op.getL1cacheAttr() ? op.getL1cacheAttr().getValue() @@ -11254,7 +12641,9 @@ class ConvertPtoStgOp final : public OpConversionPattern { FailureOr calleeName = buildL1CacheStoreCallee( op.getContext(), op.getValue().getType(), l1cache); if (failed(calleeName)) + { return rewriter.notifyMatchFailure(op, "unsupported stg signature"); + } pto::StL2Cache mode = op.getL2cacheAttr() ? op.getL2cacheAttr().getValue() @@ -11288,10 +12677,14 @@ class ConvertVPTOTypedCarrierOp final : public ConversionPattern { matchAndRewrite(Operation *op, ArrayRef operands, ConversionPatternRewriter &rewriter) const override { if (isa(op)) + { return failure(); + } Type propertyType; if (auto allocaOp = dyn_cast(op)) + { propertyType = allocaOp.getElemType(); + } else if (auto gepOp = dyn_cast(op)) propertyType = gepOp.getElemType(); if (!hasVPTOConvertibleType(op->getOperandTypes()) && @@ -11319,9 +12712,13 @@ class ConvertVPTOTypedCarrierOp final : public ConversionPattern { return rewriter.notifyMatchFailure( op, "failed to convert LLVM element type"); if (auto allocaOp = dyn_cast(converted)) + { allocaOp.setElemType(convertedPropertyType); + } else + { cast(converted).setElemType(convertedPropertyType); + } } rewriter.replaceOp(op, converted->getResults()); return success(); @@ -11788,7 +13185,9 @@ static void foldVPTOTypeCasts(ModuleOp module, TypeConverter &typeConverter) { SmallVector castsToFold; module.walk([&](UnrealizedConversionCastOp castOp) { if (castOp->getNumOperands() != 1 || castOp->getNumResults() != 1) + { return; + } if (!hasVPTOConvertibleType(castOp->getOperandTypes()) && !hasVPTOConvertibleType(castOp->getResultTypes())) return; @@ -11822,7 +13221,9 @@ static LogicalResult lowerVPTOOps(ModuleOp module, return failure(); } if (failed(materializeDecls(module, state.plannedDecls, diagOS))) + { return failure(); + } return success(); } @@ -11888,7 +13289,9 @@ static LogicalResult lowerVPTOTypes(ModuleOp module, llvm::raw_ostream &diagOS) return failure(); } if (failed(materializeDecls(module, state.plannedDecls, diagOS))) + { return failure(); + } foldVPTOTypeCasts(module, typeConverter); return success(); } @@ -11919,17 +13322,23 @@ static void normalizeFuncSignaturesForOfficialLLVMLowering(ModuleOp module) { } if (!changed) + { continue; + } auto newType = builder.getFunctionType(newInputs, newResults); funcOp.setFunctionTypeAttr(TypeAttr::get(newType)); if (funcOp.isExternal()) + { continue; + } Block &entry = funcOp.getBody().front(); for (auto [arg, newType] : llvm::zip(entry.getArguments(), newInputs)) if (arg.getType() != newType) + { arg.setType(newType); + } } } @@ -11938,7 +13347,9 @@ static void forceV300CtrlModeForVPTOFuncs(ModuleOp module) { for (func::FuncOp funcOp : module.getOps()) { if (!needsV300CtrlModeForVPTOFunc(funcOp)) + { continue; + } Block &entry = funcOp.getBody().front(); builder.setInsertionPointToStart(&entry); @@ -11959,7 +13370,9 @@ static std::optional getKernelKind(ModuleOp module) { auto kernelKind = module->getAttrOfType( FunctionKernelKindAttr::name); if (!kernelKind) + { return std::nullopt; + } return kernelKind.getKernelKind(); } @@ -11997,13 +13410,17 @@ makeDeviceEmissionOptions(const VPTOEmissionOptions &baseOptions, options.aicoreArch = options.march; options.defaultTargetCPU = options.march; if (options.march == "dav-c220-vec") + { options.defaultTargetFeatures = kC220VecTargetFeatures.str(); + } else if (options.march == "dav-c220-cube") options.defaultTargetFeatures = kC220CubeTargetFeatures.str(); else if (kind == FunctionKernelKind::Cube) options.defaultTargetFeatures = kCubeTargetFeatures.str(); else + { options.defaultTargetFeatures = kVecTargetFeatures.str(); + } } return options; } @@ -12015,9 +13432,13 @@ getUniqueDeviceModuleByKernelKind(ModuleOp module, FunctionKernelKind kind, for (ModuleOp child : module.getOps()) { auto kernelKind = getKernelKind(child); if (!kernelKind) + { continue; + } if (*kernelKind != kind) + { continue; + } if (matched) { diagOS << "VPTO LLVM emission failed: duplicate device module with " << FunctionKernelKindAttr::name << "\n"; @@ -12036,15 +13457,21 @@ static void mergeDeviceModulesByKernelKind(ModuleOp module) { for (ModuleOp child : module.getOps()) { auto kernelKind = getKernelKind(child); if (!kernelKind) + { continue; + } ModuleOp *target = nullptr; if (*kernelKind == FunctionKernelKind::Vector) + { target = &vectorModule; + } else if (*kernelKind == FunctionKernelKind::Cube) target = &cubeModule; else + { continue; + } if (!*target) { *target = child; @@ -12061,7 +13488,9 @@ static void mergeDeviceModulesByKernelKind(ModuleOp module) { } for (ModuleOp child : modulesToErase) + { child.erase(); + } } static LogicalResult renameKernelFunctionsForKernelKind(ModuleOp module, @@ -12075,7 +13504,9 @@ static LogicalResult renameKernelFunctionsForKernelKind(ModuleOp module, StringRef suffix; if (*kernelKind == FunctionKernelKind::Vector) + { suffix = kVectorSuffix; + } else if (*kernelKind == FunctionKernelKind::Cube) suffix = kCubeSuffix; else { @@ -12086,9 +13517,13 @@ static LogicalResult renameKernelFunctionsForKernelKind(ModuleOp module, for (func::FuncOp funcOp : module.getOps()) { if (!pto::hasExplicitPTOEntryAttr(funcOp)) + { continue; + } if (funcOp.getSymName().ends_with(suffix)) + { continue; + } funcOp.setSymName((funcOp.getSymName() + suffix).str()); } return success(); @@ -12113,13 +13548,19 @@ struct LowerVPTOOpsPass final SmallVector deadAllocs; getOperation().walk([&](pto::AllocTileOp alloc) { if (alloc.use_empty()) + { deadAllocs.push_back(alloc); + } }); for (pto::AllocTileOp alloc : llvm::reverse(deadAllocs)) + { alloc.erase(); + } } if (failed(lowerVPTOOps(getOperation(), march, llvm::errs()))) + { signalPassFailure(); + } } private: @@ -12132,7 +13573,9 @@ struct LowerVPTOTypesPass final void runOnOperation() override { if (failed(lowerVPTOTypes(getOperation(), llvm::errs()))) + { signalPassFailure(); + } } }; @@ -12156,7 +13599,9 @@ struct PrepareVPTOLLVMLoweringPass final pto::annotatePTOEntryFunctions(module); forceV300CtrlModeForVPTOFuncs(module); if (failed(renameKernelFunctionsForKernelKind(module, llvm::errs()))) + { signalPassFailure(); + } } }; @@ -12165,7 +13610,9 @@ collectSimtEntryFunctionNames(ModuleOp module) { llvm::StringSet simtEntries; module.walk([&](func::FuncOp funcOp) { if (funcOp->hasAttr(pto::kPTOSimtEntryAttrName)) + { simtEntries.insert(funcOp.getSymName()); + } }); return simtEntries; } @@ -12175,7 +13622,9 @@ static void applyArtifactVisibilityLinkage(ModuleOp sourceModule, llvm::StringMap externalByName; sourceModule.walk([&](func::FuncOp funcOp) { if (funcOp.isDeclaration()) + { return; + } externalByName[funcOp.getSymName()] = pto::hasExternalArtifactVisibility(funcOp); }); @@ -12183,7 +13632,9 @@ static void applyArtifactVisibilityLinkage(ModuleOp sourceModule, for (llvm::Function &function : llvmModule) { auto it = externalByName.find(function.getName()); if (it == externalByName.end()) + { continue; + } if (it->second) { function.setLinkage(llvm::GlobalValue::ExternalLinkage); continue; @@ -12215,10 +13666,14 @@ static void applySimtEntryCallingConvention( for (llvm::Instruction &inst : block) { auto *call = llvm::dyn_cast(&inst); if (!call) + { continue; + } auto *callee = call->getCalledFunction(); if (!callee || !simtEntryNames.contains(callee->getName())) + { continue; + } call->setCallingConv(llvm::CallingConv::SimtEntry); } } @@ -12233,7 +13688,9 @@ emitDeviceLLVMModule(ModuleOp deviceModule, StringRef kernelKind, if (!deviceModule) return EmittedLLVMModule{}; if (failed(applyQueriedTargetAttrs(deviceModule, options, diagOS))) + { return failure(); + } auto llvmContext = std::make_unique(); registerBuiltinDialectTranslation(*deviceModule.getContext()); @@ -12249,7 +13706,9 @@ emitDeviceLLVMModule(ModuleOp deviceModule, StringRef kernelKind, applyArtifactVisibilityLinkage(deviceModule, *llvmModule); for (llvm::Function &func : *llvmModule) { if (!func.getName().starts_with("llvm.hivm.vscatter.")) + { continue; + } // Work around a bug in older Bisheng releases: vscatter was not modeled // as writing through its destination pointer, so EarlyCSE could eliminate // a load after vscatter as redundant. @@ -12259,7 +13718,9 @@ emitDeviceLLVMModule(ModuleOp deviceModule, StringRef kernelKind, } applySimtEntryCallingConvention(*llvmModule, simtEntryNames); if (failed(attachAIVectorScopeMetadata(*llvmModule, diagOS))) + { return failure(); + } attachHIVMKernelAnnotations(*llvmModule, deviceModule); llvmModule->setModuleIdentifier(("ptoas.hivm.official." + kernelKind).str()); llvmModule->setSourceFileName(("ptoas.hivm.official." + kernelKind).str()); @@ -12327,12 +13788,16 @@ LogicalResult lowerVPTOModuleToLLVMModulesBeta1( getUniqueDeviceModuleByKernelKind( loweredModule, FunctionKernelKind::Vector, diagOS); if (failed(vectorDeviceModule)) + { return failure(); + } auto cubeDeviceModule = getUniqueDeviceModuleByKernelKind( loweredModule, FunctionKernelKind::Cube, diagOS); if (failed(cubeDeviceModule)) + { return failure(); + } if (*vectorDeviceModule) { auto vectorOptions = @@ -12341,7 +13806,9 @@ LogicalResult lowerVPTOModuleToLLVMModulesBeta1( emitDeviceLLVMModule(*vectorDeviceModule, "vector", vectorOptions, simtEntryNames, diagOS); if (failed(emitted)) + { return failure(); + } vectorModule.context = std::move(emitted->context); vectorModule.module = std::move(emitted->module); } @@ -12352,7 +13819,9 @@ LogicalResult lowerVPTOModuleToLLVMModulesBeta1( emitDeviceLLVMModule(*cubeDeviceModule, "cube", cubeOptions, simtEntryNames, diagOS); if (failed(emitted)) + { return failure(); + } cubeModule.context = std::move(emitted->context); cubeModule.module = std::move(emitted->module); } diff --git a/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp b/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp index 70c3b852eb..7d528eb9eb 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp @@ -22,8 +22,9 @@ static bool usesCANN900Lowering(const VPTOEmissionOptions &options) { static bool containsLdStDev(ModuleOp module) { bool found = false; module.walk([&](Operation *op) { - if (isa(op)) + if (isa(op)) { found = true; + } }); return found; } @@ -31,17 +32,19 @@ static bool containsLdStDev(ModuleOp module) { static LogicalResult verifyLdStDevTarget(ModuleOp module, const VPTOEmissionOptions &options, llvm::raw_ostream &diagOS) { - if (!containsLdStDev(module) || usesCANN900Lowering(options)) + if (!containsLdStDev(module) || usesCANN900Lowering(options)) { return success(); + } const bool isC220 = options.march == "dav-c220-vec" || options.march == "dav-c220-cube"; - if (isC220) + if (isC220) { diagOS << "VPTO LLVM emission failed: pto.ld_dev and pto.st_dev require " "--pto-arch=a5\n"; - else + } else { diagOS << "VPTO LLVM emission failed: pto.ld_dev and pto.st_dev require " "CANN 9.0.0 or newer official lowering\n"; + } return failure(); } @@ -49,11 +52,13 @@ LogicalResult lowerVPTOModuleToLLVMModules( ModuleOp module, const VPTOEmissionOptions &options, EmittedLLVMModule &cubeModule, EmittedLLVMModule &vectorModule, llvm::raw_ostream &diagOS) { - if (failed(verifyLdStDevTarget(module, options, diagOS))) + if (failed(verifyLdStDevTarget(module, options, diagOS))) { return failure(); - if (usesCANN900Lowering(options)) + } + if (usesCANN900Lowering(options)) { return lowerVPTOModuleToLLVMModulesCANN900(module, options, cubeModule, vectorModule, diagOS); + } return lowerVPTOModuleToLLVMModulesBeta1(module, options, cubeModule, vectorModule, diagOS); } @@ -67,8 +72,9 @@ LogicalResult lowerVPTOModuleToLLVMIRText( EmittedLLVMModule vectorModule; if (failed( lowerVPTOModuleToLLVMModules(module, options, cubeModule, vectorModule, - diagOS))) + diagOS))) { return failure(); + } llvm::raw_string_ostream os(output); bool printedAny = false; @@ -78,8 +84,9 @@ LogicalResult lowerVPTOModuleToLLVMIRText( printedAny = true; } if (cubeModule.module) { - if (printedAny) + if (printedAny) { os << "\n"; + } cubeModule.module->print(os, nullptr); os << "\n"; } @@ -96,5 +103,4 @@ LogicalResult lowerVPTOModuleToLLVMIRText( return success(); } - } // namespace mlir::pto diff --git a/lib/PTO/Transforms/VPTOLLVMEmitterHelper.cpp b/lib/PTO/Transforms/VPTOLLVMEmitterHelper.cpp index d4d4e0acba..343b48b4c2 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitterHelper.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitterHelper.cpp @@ -69,11 +69,13 @@ struct QueriedTargetAttrs { }; static bool hasPtoMemRefMemorySpace(Type type) { - if (auto memRefType = dyn_cast(type)) + if (auto memRefType = dyn_cast(type)) { return isa(memRefType.getMemorySpace()); - if (auto functionType = dyn_cast(type)) + } + if (auto functionType = dyn_cast(type)) { return llvm::any_of(functionType.getInputs(), hasPtoMemRefMemorySpace) || llvm::any_of(functionType.getResults(), hasPtoMemRefMemorySpace); + } return false; } @@ -92,16 +94,19 @@ struct ConvertPtoMemRefSpaceCarrierOp final : ConversionPattern { matchAndRewrite(Operation *op, ArrayRef operands, ConversionPatternRewriter &rewriter) const override { if (!hasPtoMemRefMemorySpace(op->getOperandTypes()) && - !hasPtoMemRefMemorySpace(op->getResultTypes())) + !hasPtoMemRefMemorySpace(op->getResultTypes())) { return failure(); - if (op->getNumRegions() != 0) + } + if (op->getNumRegions() != 0) { return rewriter.notifyMatchFailure( op, "region ops with PTO memref spaces are handled structurally"); + } FailureOr converted = convertOpResultTypes(op, operands, *typeConverter, rewriter); - if (failed(converted)) + if (failed(converted)) { return failure(); + } return success(); } }; @@ -115,8 +120,9 @@ struct ConvertMemRefReinterpretCastSpaceOp final ConversionPatternRewriter &rewriter) const override { Type convertedResultType = getTypeConverter()->convertType(op.getType()); auto memRefResultType = dyn_cast_or_null(convertedResultType); - if (!memRefResultType) + if (!memRefResultType) { return rewriter.notifyMatchFailure(op, "expected memref result type"); + } rewriter.replaceOpWithNewOp( op, memRefResultType, adaptor.getSource(), adaptor.getOffsets(), @@ -135,8 +141,9 @@ struct ConvertMemRefSubViewSpaceOp final ConversionPatternRewriter &rewriter) const override { Type convertedResultType = getTypeConverter()->convertType(op.getType()); auto memRefResultType = dyn_cast_or_null(convertedResultType); - if (!memRefResultType) + if (!memRefResultType) { return rewriter.notifyMatchFailure(op, "expected memref result type"); + } rewriter.replaceOpWithNewOp( op, memRefResultType, adaptor.getSource(), op.getMixedOffsets(), @@ -152,16 +159,19 @@ struct ConvertMemRefSpaceUnrealizedCastOp final LogicalResult matchAndRewrite(UnrealizedConversionCastOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - if (op->getNumOperands() != 1 || op->getNumResults() != 1) + if (op->getNumOperands() != 1 || op->getNumResults() != 1) { return failure(); + } if (!hasPtoMemRefMemorySpace(op->getOperandTypes()) && - !hasPtoMemRefMemorySpace(op->getResultTypes())) + !hasPtoMemRefMemorySpace(op->getResultTypes())) { return failure(); + } Type convertedResultType = getTypeConverter()->convertType(op.getResult(0).getType()); - if (!convertedResultType) + if (!convertedResultType) { return failure(); + } Value input = adaptor.getOperands().front(); if (input.getType() == convertedResultType) { @@ -174,8 +184,9 @@ struct ConvertMemRefSpaceUnrealizedCastOp final static void ensureAIVScopeDummyDecl(ModuleOp module) { SymbolTable symbolTable(module); - if (symbolTable.lookup(kAIVScopeDummyCallee)) + if (symbolTable.lookup(kAIVScopeDummyCallee)) { return; + } OpBuilder builder(module.getBodyRegion()); builder.setInsertionPointToStart(module.getBody()); @@ -187,12 +198,14 @@ static void ensureAIVScopeDummyDecl(ModuleOp module) { static bool satisfiesAIVectorScopeLatchPostcondition(llvm::Loop *loop) { llvm::BasicBlock *latch = loop->getLoopLatch(); - if (!latch) + if (!latch) { return false; + } llvm::SmallVector preds(llvm::predecessors(latch)); - if (preds.size() != 1) + if (preds.size() != 1) { return false; + } auto *predTerm = preds.front()->getTerminator(); return predTerm && predTerm->getNumSuccessors() == 1 && @@ -201,8 +214,9 @@ static bool satisfiesAIVectorScopeLatchPostcondition(llvm::Loop *loop) { static LogicalResult ensureDummyPredForAIVectorScopeLatch( llvm::Loop *loop, llvm::raw_ostream &diagOS) { - if (satisfiesAIVectorScopeLatchPostcondition(loop)) + if (satisfiesAIVectorScopeLatchPostcondition(loop)) { return success(); + } llvm::BasicBlock *latch = loop->getLoopLatch(); if (!latch) { @@ -239,12 +253,14 @@ static FailureOr extractQuotedLLVMFnAttr(llvm::StringRef ir, pattern += key.str(); pattern += "\"=\""; size_t start = ir.find(pattern); - if (start == llvm::StringRef::npos) + if (start == llvm::StringRef::npos) { return failure(); + } start += pattern.size(); size_t end = ir.find('"', start); - if (end == llvm::StringRef::npos || end <= start) + if (end == llvm::StringRef::npos || end <= start) { return failure(); + } return ir.slice(start, end).str(); } @@ -261,8 +277,9 @@ queryDefaultTargetAttrs(const VPTOEmissionOptions &options, std::string cacheKey = options.targetTriple + "|" + options.march + "|" + options.aicoreArch; - if (auto it = cache.find(cacheKey); it != cache.end()) + if (auto it = cache.find(cacheKey); it != cache.end()) { return it->second; + } auto bisheng = llvm::sys::findProgramByName("bisheng"); if (!bisheng) { @@ -332,8 +349,9 @@ queryDefaultTargetAttrs(const VPTOEmissionOptions &options, }; llvm::SmallVector args; args.reserve(argStorage.size()); - for (const std::string &arg : argStorage) + for (const std::string &arg : argStorage) { args.push_back(arg); + } std::string execErr; bool execFailed = false; @@ -349,13 +367,16 @@ queryDefaultTargetAttrs(const VPTOEmissionOptions &options, if (execFailed || rc != 0) { diagOS << "VPTO LLVM emission failed: bisheng target query failed\n"; diagOS << "Command:"; - for (llvm::StringRef arg : args) + for (llvm::StringRef arg : args) { diagOS << " " << arg; + } diagOS << "\n"; - if (!execErr.empty()) + if (!execErr.empty()) { diagOS << execErr << "\n"; - if (!stderrText.empty()) + } + if (!stderrText.empty()) { diagOS << stderrText << "\n"; + } return failure(); } @@ -393,8 +414,9 @@ void materializeVecScopeCarrierLoops(ModuleOp module) { IRRewriter rewriter(module.getContext()); for (pto::VecScopeOp vecScope : llvm::reverse(scopes)) { - if (!vecScope || vecScope.getBody().empty()) + if (!vecScope || vecScope.getBody().empty()) { continue; + } rewriter.setInsertionPoint(vecScope); auto loc = vecScope.getLoc(); @@ -421,8 +443,9 @@ void materializeVecScopeCarrierLoops(ModuleOp module) { }); for (pto::StrictVecScopeOp strictVecScope : llvm::reverse(strictScopes)) { - if (!strictVecScope || strictVecScope.getBody().empty()) + if (!strictVecScope || strictVecScope.getBody().empty()) { continue; + } rewriter.setInsertionPoint(strictVecScope); auto loc = strictVecScope.getLoc(); @@ -436,12 +459,14 @@ void materializeVecScopeCarrierLoops(ModuleOp module) { IRMapping mapping; for (auto [blockArg, capture] : - llvm::zip(strictBody.getArguments(), strictVecScope.getCaptures())) + llvm::zip(strictBody.getArguments(), strictVecScope.getCaptures())) { mapping.map(blockArg, capture); + } rewriter.setInsertionPoint(yield); - for (Operation &nested : strictBody.getOperations()) + for (Operation &nested : strictBody.getOperations()) { rewriter.clone(nested, mapping); + } rewriter.create(loc, kAIVScopeDummyCallee, TypeRange{}, ValueRange{}); @@ -452,12 +477,14 @@ void materializeVecScopeCarrierLoops(ModuleOp module) { LogicalResult attachAIVectorScopeMetadata(llvm::Module &llvmModule, llvm::raw_ostream &diagOS) { llvm::Function *dummyCallee = llvmModule.getFunction(kAIVScopeDummyCallee); - if (!dummyCallee) + if (!dummyCallee) { return success(); + } for (llvm::Function &function : llvmModule) { - if (function.isDeclaration()) + if (function.isDeclaration()) { continue; + } llvm::DominatorTree dt(function); llvm::LoopInfo loopInfo(dt); @@ -465,8 +492,9 @@ LogicalResult attachAIVectorScopeMetadata(llvm::Module &llvmModule, for (llvm::BasicBlock &block : function) { for (llvm::Instruction &inst : block) { auto *call = dyn_cast(&inst); - if (call && call->getCalledFunction() == dummyCallee) + if (call && call->getCalledFunction() == dummyCallee) { dummyCalls.push_back(call); + } } } @@ -495,8 +523,9 @@ LogicalResult attachAIVectorScopeMetadata(llvm::Module &llvmModule, } } - if (failed(ensureDummyPredForAIVectorScopeLatch(loop, diagOS))) + if (failed(ensureDummyPredForAIVectorScopeLatch(loop, diagOS))) { return failure(); + } dt.recalculate(function); loopInfo.releaseMemory(); @@ -529,18 +558,22 @@ LogicalResult attachAIVectorScopeMetadata(llvm::Module &llvmModule, } } - if (dummyCallee->use_empty()) + if (dummyCallee->use_empty()) { dummyCallee->eraseFromParent(); + } return success(); } constexpr uint32_t getSimtMaxRegistersForThreads(uint32_t maxThreads) { - if (maxThreads > 1024) + if (maxThreads > 1024) { return 16; - if (maxThreads > 512) + } + if (maxThreads > 512) { return 32; - if (maxThreads > 256) + } + if (maxThreads > 256) { return 64; + } return 128; } @@ -563,13 +596,15 @@ void attachHIVMKernelAnnotations(llvm::Module &llvmModule, ptoEntryFunctions.insert(symName); } - if (!funcOp->hasAttr(pto::kPTOSimtEntryAttrName)) + if (!funcOp->hasAttr(pto::kPTOSimtEntryAttrName)) { return; + } uint32_t maxThreads = kDefaultSimtMaxThreads; if (auto attr = - funcOp->getAttrOfType(pto::kPTOSimtMaxThreadsAttrName)) + funcOp->getAttrOfType(pto::kPTOSimtMaxThreadsAttrName)) { maxThreads = static_cast(attr.getInt()); + } simtMaxThreadsByName[symName] = maxThreads; }); @@ -578,10 +613,12 @@ void attachHIVMKernelAnnotations(llvm::Module &llvmModule, for (llvm::BasicBlock &block : function) { for (llvm::Instruction &inst : block) { auto *call = llvm::dyn_cast(&inst); - if (!call) + if (!call) { continue; - if (call->getCallingConv() == llvm::CallingConv::SimtEntry) + } + if (call->getCallingConv() == llvm::CallingConv::SimtEntry) { return true; + } } } return false; @@ -612,13 +649,15 @@ void attachHIVMKernelAnnotations(llvm::Module &llvmModule, }; for (llvm::Function &function : llvmModule) { - if (function.isDeclaration()) + if (function.isDeclaration()) { continue; + } if (function.getCallingConv() == llvm::CallingConv::SimtEntry) { uint32_t maxThreads = kDefaultSimtMaxThreads; if (auto it = simtMaxThreadsByName.find(function.getName()); - it != simtMaxThreadsByName.end()) + it != simtMaxThreadsByName.end()) { maxThreads = it->second; + } uint32_t maxRegisters = getSimtMaxRegistersForThreads(maxThreads); addLLVMFunctionI32Annotation(function, "simt-max-threads", maxThreads); @@ -628,21 +667,24 @@ void attachHIVMKernelAnnotations(llvm::Module &llvmModule, addHIVMModuleI32Annotation("simt-max-registers", maxRegisters); continue; } - if (function.getLinkage() != llvm::GlobalValue::ExternalLinkage) + if (function.getLinkage() != llvm::GlobalValue::ExternalLinkage) { continue; + } llvm::StringRef name = function.getName(); - if (!ptoEntryFunctions.contains(name)) + if (!ptoEntryFunctions.contains(name)) { continue; - if (name.contains(".extracted") || name.contains(".vector.thread")) + } + if (name.contains(".extracted") || name.contains(".vector.thread")) { continue; + } addAnnotation(function, "kernel"); addAnnotation(function, "kernel_with_simd"); - if (callsSimtEntry(function)) + if (callsSimtEntry(function)) { addAnnotation(function, "kernel_with_simt"); + } } - } LogicalResult @@ -651,8 +693,9 @@ applyQueriedTargetAttrs(ModuleOp module, const VPTOEmissionOptions &options, FailureOr attrs = queryDefaultTargetAttrs(options, diagOS); if (failed(attrs)) { if (options.defaultTargetCPU.empty() || - options.defaultTargetFeatures.empty()) + options.defaultTargetFeatures.empty()) { return failure(); + } diagOS << "VPTO LLVM emission: falling back to configured default target " "attributes\n"; attrs = QueriedTargetAttrs{options.defaultTargetCPU, diff --git a/lib/PTO/Transforms/VPTOMaskSimplify.cpp b/lib/PTO/Transforms/VPTOMaskSimplify.cpp index 90827555db..eff3273b1e 100644 --- a/lib/PTO/Transforms/VPTOMaskSimplify.cpp +++ b/lib/PTO/Transforms/VPTOMaskSimplify.cpp @@ -43,8 +43,9 @@ struct SimplifyAllTruePredicateReorder : public OpRewritePattern { LogicalResult matchAndRewrite(OpTy op, PatternRewriter &rewriter) const override { - if (!isAllTrueMask(op.getLhs()) || !isAllTrueMask(op.getRhs())) + if (!isAllTrueMask(op.getLhs()) || !isAllTrueMask(op.getRhs())) { return failure(); + } rewriter.replaceOp(op, {op.getLhs(), op.getRhs()}); return success(); @@ -62,8 +63,9 @@ struct VPTOMaskSimplifyPass SimplifyAllTruePredicateReorder, SimplifyAllTruePredicateReorder>(&getContext()); - if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) + if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) { signalPassFailure(); + } } }; diff --git a/lib/PTO/Transforms/VPTONormalizeContainer.cpp b/lib/PTO/Transforms/VPTONormalizeContainer.cpp index d5b22af3d6..2ea6b606e2 100644 --- a/lib/PTO/Transforms/VPTONormalizeContainer.cpp +++ b/lib/PTO/Transforms/VPTONormalizeContainer.cpp @@ -44,8 +44,9 @@ static LogicalResult verifyNormalizedVPTOContainer(ModuleOp module) { } } - if (hasChildModules) + if (hasChildModules) { return success(); + } return module.emitError() << "expected VPTO input to be a kernel submodule with " @@ -60,10 +61,12 @@ struct VPTONormalizeContainerPass if (isVPTOKernelSubmodule(module)) { MLIRContext *context = module.getContext(); SmallVector outerAttrs; - for (NamedAttribute attr : module->getAttrs()) + for (NamedAttribute attr : module->getAttrs()) { if (attr.getName() != SymbolTable::getSymbolAttrName() && - attr.getName() != FunctionKernelKindAttr::name) + attr.getName() != FunctionKernelKindAttr::name) { outerAttrs.push_back(attr); + } + } auto child = ModuleOp::create(module.getLoc()); child->setAttrs(module->getAttrDictionary()); @@ -74,8 +77,9 @@ struct VPTONormalizeContainerPass module.getBodyRegion().front().push_back(child.getOperation()); } - if (failed(verifyNormalizedVPTOContainer(module))) + if (failed(verifyNormalizedVPTOContainer(module))) { signalPassFailure(); + } } }; diff --git a/lib/PTO/Transforms/VPTOOptimizeVcvt.cpp b/lib/PTO/Transforms/VPTOOptimizeVcvt.cpp index 9a714e73c8..c59942cace 100644 --- a/lib/PTO/Transforms/VPTOOptimizeVcvt.cpp +++ b/lib/PTO/Transforms/VPTOOptimizeVcvt.cpp @@ -31,12 +31,15 @@ static bool isOddPart(StringRef part) { } static bool isAllTrueMask(Value mask) { - if (auto op = mask.getDefiningOp()) + if (auto op = mask.getDefiningOp()) { return op.getPattern() == "PAT_ALL"; - if (auto op = mask.getDefiningOp()) + } + if (auto op = mask.getDefiningOp()) { return op.getPattern() == "PAT_ALL"; - if (auto op = mask.getDefiningOp()) + } + if (auto op = mask.getDefiningOp()) { return op.getPattern() == "PAT_ALL"; + } return false; } @@ -47,12 +50,14 @@ static bool isPairEquivalentLoadDist(StringRef dist) { } static bool hasEvenOddEquivalentLanes(Value value) { - if (value.getDefiningOp()) + if (value.getDefiningOp()) { return true; + } auto load = value.getDefiningOp(); - if (!load || value != load.getResult()) + if (!load || value != load.getResult()) { return false; + } std::optional dist = load.getDist(); return dist && isPairEquivalentLoadDist(*dist); @@ -61,8 +66,9 @@ static bool hasEvenOddEquivalentLanes(Value value) { static bool isNarrowToWideVcvt(VcvtOp op) { auto inputType = dyn_cast(op.getInput().getType()); auto resultType = dyn_cast(op.getResult().getType()); - if (!inputType || !resultType) + if (!inputType || !resultType) { return false; + } unsigned inputBits = getPTOStorageElemBitWidth(inputType.getElementType()); unsigned resultBits = getPTOStorageElemBitWidth(resultType.getElementType()); @@ -70,8 +76,9 @@ static bool isNarrowToWideVcvt(VcvtOp op) { } static Value stripVbitcasts(Value value) { - while (auto bitcast = value.getDefiningOp()) + while (auto bitcast = value.getDefiningOp()) { value = bitcast.getInput(); + } return value; } @@ -82,35 +89,42 @@ struct AlignedUnsignedWidening { static bool isZeroGapLoad(Value value, AlignedUnsignedWidening widening) { auto load = stripVbitcasts(value).getDefiningOp(); - if (!load) + if (!load) { return false; + } auto loadType = dyn_cast(load.getResult().getType()); std::optional dist = load.getDist(); if (!loadType || !dist || getPTOStorageElemBitWidth(loadType.getElementType()) != - widening.payloadBits) + widening.payloadBits) { return false; + } - if (widening.payloadBits == 8 && widening.carrierBits == 16) + if (widening.payloadBits == 8 && widening.carrierBits == 16) { return *dist == "UNPK_B8"; - if (widening.payloadBits == 16 && widening.carrierBits == 32) + } + if (widening.payloadBits == 16 && widening.carrierBits == 32) { return *dist == "UNPK_B16"; - if (widening.payloadBits == 8 && widening.carrierBits == 32) + } + if (widening.payloadBits == 8 && widening.carrierBits == 32) { return *dist == "UNPK4"; + } return false; } static bool isZeroGapNarrowingVcvt(Value value, AlignedUnsignedWidening widening) { auto cvt = value.getDefiningOp(); - if (!cvt || !isAllTrueMask(cvt.getMask())) + if (!cvt || !isAllTrueMask(cvt.getMask())) { return false; + } auto inputType = dyn_cast(cvt.getInput().getType()); auto resultType = dyn_cast(cvt.getResult().getType()); - if (!inputType || !resultType) + if (!inputType || !resultType) { return false; + } unsigned inputBits = getPTOStorageElemBitWidth(inputType.getElementType()); unsigned resultBits = getPTOStorageElemBitWidth(resultType.getElementType()); @@ -118,12 +132,14 @@ static bool isZeroGapNarrowingVcvt(Value value, resultBits != widening.payloadBits || inputBits <= resultBits || inputType.getElementCount() * inputBits != - resultType.getElementCount() * resultBits) + resultType.getElementCount() * resultBits) { return false; + } std::optional part = cvt.getPart(); - if (!part) + if (!part) { return false; + } return (inputBits == resultBits * 2 && *part == "EVEN") || (inputBits == resultBits * 4 && *part == "P0"); } @@ -134,29 +150,35 @@ static bool isZeroGapNarrowingVcvt(Value value, static bool isCanonicalZeroGapCarrier(Value value, AlignedUnsignedWidening widening) { value = stripVbitcasts(value); - if (isZeroGapLoad(value, widening)) + if (isZeroGapLoad(value, widening)) { return true; - if (isZeroGapNarrowingVcvt(value, widening)) + } + if (isZeroGapNarrowingVcvt(value, widening)) { return true; + } auto unpack = value.getDefiningOp(); - if (!unpack) + if (!unpack) { return false; + } auto sourceType = dyn_cast(unpack.getSrc().getType()); auto resultType = dyn_cast(unpack.getResult().getType()); - if (!sourceType || !resultType) + if (!sourceType || !resultType) { return false; + } unsigned sourceBits = getPTOStorageElemBitWidth(sourceType.getElementType()); unsigned resultBits = getPTOStorageElemBitWidth(resultType.getElementType()); if (sourceBits == 0 || resultBits != widening.carrierBits || - resultBits != sourceBits * 2 || widening.payloadBits > sourceBits) + resultBits != sourceBits * 2 || widening.payloadBits > sourceBits) { return false; + } - if (widening.payloadBits == sourceBits) + if (widening.payloadBits == sourceBits) { return true; + } return isCanonicalZeroGapCarrier( unpack.getSrc(), {widening.payloadBits, sourceBits}); } @@ -166,29 +188,34 @@ matchAlignedUnsignedWidening(VcvtOp op) { auto inputType = dyn_cast(op.getInput().getType()); auto resultType = dyn_cast(op.getResult().getType()); if (!inputType || !resultType || op.getRndAttr() || op.getSatAttr() || - !isAllTrueMask(op.getMask())) + !isAllTrueMask(op.getMask())) { return std::nullopt; + } auto inputElementType = dyn_cast(inputType.getElementType()); auto resultElementType = dyn_cast(resultType.getElementType()); if (!inputElementType || !resultElementType || - !inputElementType.isUnsigned() || !resultElementType.isUnsigned()) + !inputElementType.isUnsigned() || !resultElementType.isUnsigned()) { return std::nullopt; + } unsigned inputBits = inputElementType.getWidth(); unsigned resultBits = resultElementType.getWidth(); if (inputBits >= resultBits || inputType.getElementCount() * inputBits != - resultType.getElementCount() * resultBits) + resultType.getElementCount() * resultBits) { return std::nullopt; + } std::optional part = op.getPart(); - if (!part) + if (!part) { return std::nullopt; + } if ((resultBits == inputBits * 2 && *part != "EVEN") || (resultBits == inputBits * 4 && *part != "P0") || - (resultBits != inputBits * 2 && resultBits != inputBits * 4)) + (resultBits != inputBits * 2 && resultBits != inputBits * 4)) { return std::nullopt; + } return AlignedUnsignedWidening{inputBits, resultBits}; } @@ -201,8 +228,9 @@ struct CanonicalizeEquivalentPartPattern : public OpRewritePattern { std::optional part = op.getPart(); if (!part || !isOddPart(*part) || !isNarrowToWideVcvt(op) || !isAllTrueMask(op.getMask()) || - !hasEvenOddEquivalentLanes(op.getInput())) + !hasEvenOddEquivalentLanes(op.getInput())) { return failure(); + } rewriter.modifyOpInPlace( op, [&] { op.setPartAttr(rewriter.getStringAttr("EVEN")); }); @@ -217,8 +245,9 @@ struct FoldZeroGapExtensionPattern : public OpRewritePattern { PatternRewriter &rewriter) const override { std::optional widening = matchAlignedUnsignedWidening(op); - if (!widening || !isCanonicalZeroGapCarrier(op.getInput(), *widening)) + if (!widening || !isCanonicalZeroGapCarrier(op.getInput(), *widening)) { return failure(); + } Value carrier = stripVbitcasts(op.getInput()); Value result = carrier.getType() == op.getResult().getType() @@ -239,8 +268,9 @@ struct VPTOOptimizeVcvtPass RewritePatternSet patterns(&getContext()); patterns.add(&getContext()); - if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) + if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) { signalPassFailure(); + } } }; diff --git a/lib/PTO/Transforms/VPTOPtrCastCleanup.cpp b/lib/PTO/Transforms/VPTOPtrCastCleanup.cpp index fc238fac5f..9c27752a60 100644 --- a/lib/PTO/Transforms/VPTOPtrCastCleanup.cpp +++ b/lib/PTO/Transforms/VPTOPtrCastCleanup.cpp @@ -31,32 +31,39 @@ struct CollapsePtrMemRefPtrBridgePattern LogicalResult matchAndRewrite(UnrealizedConversionCastOp op, PatternRewriter &rewriter) const override { - if (op->getNumOperands() != 1 || op->getNumResults() != 1) + if (op->getNumOperands() != 1 || op->getNumResults() != 1) { return failure(); + } auto resultPtrType = dyn_cast(op.getResult(0).getType()); - if (!resultPtrType) + if (!resultPtrType) { return failure(); + } auto castOp = op.getOperand(0).getDefiningOp(); - if (!castOp || castOp->getNumOperands() != 1) + if (!castOp || castOp->getNumOperands() != 1) { return failure(); + } auto innerCast = castOp.getSource().getDefiningOp(); if (!innerCast || innerCast->getNumOperands() != 1 || - innerCast->getNumResults() != 1) + innerCast->getNumResults() != 1) { return failure(); + } Value basePtr = innerCast.getOperand(0); - if (basePtr.getType() != resultPtrType) + if (basePtr.getType() != resultPtrType) { return failure(); + } rewriter.replaceOp(op, basePtr); - if (castOp->use_empty()) + if (castOp->use_empty()) { rewriter.eraseOp(castOp); - if (innerCast->use_empty()) + } + if (innerCast->use_empty()) { rewriter.eraseOp(innerCast); + } return success(); } }; @@ -69,8 +76,9 @@ struct VPTOPtrCastCleanupPass void runOnOperation() override { RewritePatternSet patterns(&getContext()); patterns.add(&getContext()); - if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) + if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) { signalPassFailure(); + } } }; diff --git a/lib/PTO/Transforms/VPTOPtrNormalize.cpp b/lib/PTO/Transforms/VPTOPtrNormalize.cpp index 834a521d31..97c4d377a2 100644 --- a/lib/PTO/Transforms/VPTOPtrNormalize.cpp +++ b/lib/PTO/Transforms/VPTOPtrNormalize.cpp @@ -39,39 +39,45 @@ namespace { static pto::AddressSpaceAttr getPointerMemorySpace(Attribute memorySpace, MLIRContext *ctx) { - if (auto addrSpace = dyn_cast_or_null(memorySpace)) + if (auto addrSpace = dyn_cast_or_null(memorySpace)) { return addrSpace; - if (auto intAttr = dyn_cast_or_null(memorySpace)) + } + if (auto intAttr = dyn_cast_or_null(memorySpace)) { return pto::AddressSpaceAttr::get( ctx, static_cast(intAttr.getInt())); + } return {}; } static bool needsSubviewPtrConversion(memref::SubViewOp op) { auto resultType = dyn_cast(op.getType()); - if (!resultType) + if (!resultType) { return false; + } return static_cast( getPointerMemorySpace(resultType.getMemorySpace(), op.getContext())); } static Type convertSubviewResultType(Type type) { auto memrefType = dyn_cast(type); - if (!memrefType) + if (!memrefType) { return type; + } auto memorySpace = getPointerMemorySpace(memrefType.getMemorySpace(), type.getContext()); - if (!memorySpace) + if (!memorySpace) { return type; + } return pto::PtrType::get(type.getContext(), memrefType.getElementType(), memorySpace); } static bool hasPtrNormalizeConvertibleType(Type type) { - if (isa(type)) + if (isa(type)) { return true; + } auto memrefType = dyn_cast(type); return memrefType && static_cast(getPointerMemorySpace( memrefType.getMemorySpace(), type.getContext())); @@ -94,8 +100,9 @@ static bool hasPtrNormalizeMemRefType(TypeRange types) { } static bool isTransientPtrMemRefBridge(Value value) { - if (!isa(value.getType())) + if (!isa(value.getType())) { return false; + } auto cast = value.getDefiningOp(); return cast && cast->getNumOperands() == 1 && cast->getNumResults() == 1 && isa(cast.getOperand(0).getType()); @@ -104,8 +111,9 @@ static bool isTransientPtrMemRefBridge(Value value) { static FailureOr> convertTypes(const TypeConverter &typeConverter, TypeRange types) { SmallVector convertedTypes; - if (failed(typeConverter.convertTypes(types, convertedTypes))) + if (failed(typeConverter.convertTypes(types, convertedTypes))) { return failure(); + } return convertedTypes; } @@ -113,8 +121,9 @@ static bool isMemRefType(Type type) { return isa(type); } static Value materializeUnrealizedCast(OpBuilder &builder, Type resultType, ValueRange inputs, Location loc) { - if (inputs.size() != 1) + if (inputs.size() != 1) { return {}; + } return builder .create(loc, TypeRange{resultType}, inputs) .getResult(0); @@ -124,44 +133,52 @@ static LogicalResult computeSubviewElementOffset(memref::SubViewOp op, PatternRewriter &rewriter, Value &offset) { auto sourceType = dyn_cast(op.getSource().getType()); - if (!sourceType) + if (!sourceType) { return failure(); + } SmallVector strides; int64_t baseOffset = 0; if (failed(mlir::pto::getPTOMemRefStridesAndOffset(sourceType, strides, - baseOffset))) + baseOffset))) { return failure(); + } // The SSA source already names the base address after ptr-boundary // normalization. A dynamic memref layout offset here is metadata we can // ignore for ptr normalization and model as zero. - if (baseOffset == ShapedType::kDynamic) + if (baseOffset == ShapedType::kDynamic) { baseOffset = 0; + } Location loc = op.getLoc(); Value total = rewriter.create(loc, baseOffset); ArrayRef staticOffsets = op.getStaticOffsets(); ValueRange dynamicOffsets = op.getOffsets(); - if (staticOffsets.size() != strides.size()) + if (staticOffsets.size() != strides.size()) { return failure(); + } unsigned dynamicIndex = 0; for (auto [staticOffset, stride] : llvm::zip(staticOffsets, strides)) { - if (stride == 0) + if (stride == 0) { continue; - if (stride == ShapedType::kDynamic) + } + if (stride == ShapedType::kDynamic) { return failure(); + } Value idx; if (ShapedType::isDynamic(staticOffset)) { - if (dynamicIndex >= dynamicOffsets.size()) + if (dynamicIndex >= dynamicOffsets.size()) { return failure(); + } idx = dynamicOffsets[dynamicIndex++]; } else { idx = rewriter.create(loc, staticOffset); } - if (!idx.getType().isIndex()) + if (!idx.getType().isIndex()) { return failure(); + } if (stride != 1) { Value strideValue = @@ -170,8 +187,9 @@ static LogicalResult computeSubviewElementOffset(memref::SubViewOp op, } total = rewriter.create(loc, total, idx); } - if (dynamicIndex != dynamicOffsets.size()) + if (dynamicIndex != dynamicOffsets.size()) { return failure(); + } offset = total; return success(); @@ -179,19 +197,23 @@ static LogicalResult computeSubviewElementOffset(memref::SubViewOp op, static Value materializeSubviewInputPtr(Value source, PatternRewriter &rewriter, Location loc) { - if (!source) + if (!source) { return {}; - if (isa(source.getType())) + } + if (isa(source.getType())) { return source; + } auto memrefType = dyn_cast(source.getType()); - if (!memrefType) + if (!memrefType) { return {}; + } auto memorySpace = getPointerMemorySpace(memrefType.getMemorySpace(), rewriter.getContext()); - if (!memorySpace) + if (!memorySpace) { return {}; + } auto ptrType = pto::PtrType::get(rewriter.getContext(), memrefType.getElementType(), memorySpace); @@ -200,78 +222,96 @@ static Value materializeSubviewInputPtr(Value source, PatternRewriter &rewriter, static Value materializeScalarAccessPtr(Value source, PatternRewriter &rewriter, Location loc) { - if (!source) + if (!source) { return {}; + } if (auto cast = source.getDefiningOp()) { - if (cast->getNumOperands() != 1 || cast->getNumResults() != 1) + if (cast->getNumOperands() != 1 || cast->getNumResults() != 1) { return {}; + } Value input = cast.getOperands().front(); Value ptr = materializeScalarAccessPtr(input, rewriter, loc); - if (!ptr) + if (!ptr) { return {}; + } auto resultType = dyn_cast(source.getType()); - if (!resultType) + if (!resultType) { return ptr; - if (ptr.getType() == resultType) + } + if (ptr.getType() == resultType) { return ptr; + } return rewriter.create(loc, resultType, ptr); } - if (isa(source.getType())) + if (isa(source.getType())) { return source; + } - if (auto cast = source.getDefiningOp()) + if (auto cast = source.getDefiningOp()) { return materializeScalarAccessPtr(cast.getSource(), rewriter, loc); + } if (auto reinterpret = source.getDefiningOp()) { auto ptrType = dyn_cast(convertSubviewResultType(source.getType())); - if (!ptrType) + if (!ptrType) { return {}; + } Value basePtr = materializeScalarAccessPtr(reinterpret.getSource(), rewriter, loc); - if (!basePtr) + if (!basePtr) { return {}; - if (basePtr.getType() != ptrType) + } + if (basePtr.getType() != ptrType) { basePtr = rewriter.create(loc, ptrType, basePtr); + } ArrayRef staticOffsets = reinterpret.getStaticOffsets(); - if (staticOffsets.size() != 1) + if (staticOffsets.size() != 1) { return {}; + } int64_t staticOffset = staticOffsets.front(); if (!ShapedType::isDynamic(staticOffset)) { - if (staticOffset == 0) + if (staticOffset == 0) { return basePtr; + } Value offset = rewriter.create(loc, staticOffset); return rewriter.create(loc, ptrType, basePtr, offset); } ValueRange dynamicOffsets = reinterpret.getOffsets(); - if (dynamicOffsets.size() != 1 || !dynamicOffsets.front().getType().isIndex()) + if (dynamicOffsets.size() != 1 || !dynamicOffsets.front().getType().isIndex()) { return {}; + } return rewriter.create(loc, ptrType, basePtr, dynamicOffsets.front()); } if (auto subview = source.getDefiningOp()) { - if (!needsSubviewPtrConversion(subview)) + if (!needsSubviewPtrConversion(subview)) { return {}; + } Value basePtr = materializeScalarAccessPtr(subview.getSource(), rewriter, loc); - if (!basePtr) + if (!basePtr) { return {}; + } Value offset; - if (failed(computeSubviewElementOffset(subview, rewriter, offset))) + if (failed(computeSubviewElementOffset(subview, rewriter, offset))) { return {}; + } auto ptrType = dyn_cast(convertSubviewResultType(source.getType())); - if (!ptrType) + if (!ptrType) { return {}; - if (basePtr.getType() != ptrType) + } + if (basePtr.getType() != ptrType) { basePtr = rewriter.create(loc, ptrType, basePtr); + } return rewriter.create(loc, ptrType, basePtr, offset); } @@ -284,10 +324,12 @@ static Value materializeScalarAccessPtr(Value source, PatternRewriter &rewriter, static Value materializeBoundaryOperandPtr(Value source, PatternRewriter &rewriter, Location loc) { - if (!source) + if (!source) { return {}; - if (isa(source.getType())) + } + if (isa(source.getType())) { return source; + } return materializeScalarAccessPtr(source, rewriter, loc); } @@ -297,21 +339,25 @@ static LogicalResult rewriteBufferLikeBoundaryOp( StringRef sourceRole, StringRef destinationRole) { Value source = materializeBoundaryOperandPtr(adaptor.getOperands()[0], rewriter, op.getLoc()); - if (!source) + if (!source) { return rewriter.notifyMatchFailure( op, (Twine("failed to materialize ") + sourceRole + " ptr").str()); - if (!isa(source.getType())) + } + if (!isa(source.getType())) { return rewriter.notifyMatchFailure( op, (Twine("expected ptr-form ") + sourceRole).str()); + } Value destination = materializeBoundaryOperandPtr(adaptor.getOperands()[1], rewriter, op.getLoc()); - if (!destination) + if (!destination) { return rewriter.notifyMatchFailure( op, (Twine("failed to materialize ") + destinationRole + " ptr").str()); - if (!isa(destination.getType())) + } + if (!isa(destination.getType())) { return rewriter.notifyMatchFailure( op, (Twine("expected ptr-form ") + destinationRole).str()); + } SmallVector operands(adaptor.getOperands().begin(), adaptor.getOperands().end()); @@ -336,8 +382,9 @@ struct ConvertTileBufAddrToPtrPattern matchAndRewrite(pto::TileBufAddrOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { Type convertedType = getTypeConverter()->convertType(op.getDst().getType()); - if (!isa(convertedType)) + if (!isa(convertedType)) { return failure(); + } rewriter.replaceOpWithNewOp(op, convertedType, adaptor.getSrc()); @@ -354,8 +401,9 @@ struct ConvertIntToPtrToCastPtrPattern ConversionPatternRewriter &rewriter) const override { Type convertedType = getTypeConverter()->convertType(op.getResult().getType()); - if (!isa(convertedType)) + if (!isa(convertedType)) { return rewriter.notifyMatchFailure(op, "expected pointer result type"); + } rewriter.replaceOpWithNewOp(op, convertedType, adaptor.getAddr()); @@ -371,10 +419,12 @@ struct ConvertPtrToIntToCastPtrPattern matchAndRewrite(pto::PtrToIntOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { Type convertedType = getTypeConverter()->convertType(op.getResult().getType()); - if (!isa(convertedType)) + if (!isa(convertedType)) { return rewriter.notifyMatchFailure(op, "expected integer result type"); - if (!isa(adaptor.getPtr().getType())) + } + if (!isa(adaptor.getPtr().getType())) { return rewriter.notifyMatchFailure(op, "expected pointer input type"); + } rewriter.replaceOpWithNewOp(op, convertedType, adaptor.getPtr()); @@ -390,19 +440,22 @@ struct ConvertCastPtrPattern : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Type convertedResultType = getTypeConverter()->convertType(op.getResult().getType()); - if (!convertedResultType) + if (!convertedResultType) { return failure(); + } Value input = adaptor.getInput(); Type inputType = input.getType(); - if (isMemRefType(inputType) || isMemRefType(convertedResultType)) + if (isMemRefType(inputType) || isMemRefType(convertedResultType)) { return rewriter.notifyMatchFailure(op, "memref castptr must be eliminated"); + } if (!isa(inputType) || - !isa(convertedResultType)) + !isa(convertedResultType)) { return rewriter.notifyMatchFailure(op, "expected ptr/int castptr operands"); + } if (inputType == convertedResultType) { rewriter.replaceOp(op, input); @@ -421,24 +474,28 @@ struct ConvertSubviewToAddPtrPattern LogicalResult matchAndRewrite(memref::SubViewOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - if (!needsSubviewPtrConversion(op)) + if (!needsSubviewPtrConversion(op)) { return failure(); + } auto ptrType = dyn_cast(getTypeConverter()->convertType(op.getType())); - if (!ptrType) + if (!ptrType) { return rewriter.notifyMatchFailure(op, "expected ptr result type"); + } Value basePtr = materializeSubviewInputPtr(adaptor.getSource(), rewriter, op.getLoc()); - if (!basePtr) + if (!basePtr) { return rewriter.notifyMatchFailure(op, "failed to materialize subview input ptr"); + } Value offset; - if (failed(computeSubviewElementOffset(op, rewriter, offset))) + if (failed(computeSubviewElementOffset(op, rewriter, offset))) { return rewriter.notifyMatchFailure(op, "failed to compute subview element offset"); + } rewriter.replaceOpWithNewOp(op, ptrType, basePtr, offset); return success(); @@ -451,15 +508,17 @@ struct ConvertVldsSubviewOperandPattern : public OpConversionPattern(adaptor.getSource().getType())) + if (!isa(adaptor.getSource().getType())) { return failure(); + } OperationState state(op.getLoc(), op->getName().getStringRef()); state.addOperands({adaptor.getSource(), adaptor.getOffset()}); FailureOr> resultTypes = convertTypes(*getTypeConverter(), op->getResultTypes()); - if (failed(resultTypes)) + if (failed(resultTypes)) { return failure(); + } state.addTypes(*resultTypes); state.addAttributes(op->getAttrs()); Operation *newOp = rewriter.create(state); @@ -474,8 +533,9 @@ struct ConvertVstsSubviewOperandPattern : public OpConversionPattern(adaptor.getDestination().getType())) + if (!isa(adaptor.getDestination().getType())) { return failure(); + } OperationState state(op.getLoc(), op->getName().getStringRef()); state.addOperands( @@ -483,8 +543,9 @@ struct ConvertVstsSubviewOperandPattern : public OpConversionPattern> resultTypes = convertTypes(*getTypeConverter(), op->getResultTypes()); - if (failed(resultTypes)) + if (failed(resultTypes)) { return failure(); + } state.addTypes(*resultTypes); state.addAttributes(op->getAttrs()); Operation *newOp = rewriter.create(state); @@ -500,8 +561,9 @@ struct ConvertVsstbSubviewOperandPattern LogicalResult matchAndRewrite(pto::VsstbOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - if (!isa(adaptor.getDestination().getType())) + if (!isa(adaptor.getDestination().getType())) { return failure(); + } OperationState state(op.getLoc(), op->getName().getStringRef()); state.addOperands({adaptor.getValue(), adaptor.getDestination(), @@ -509,8 +571,9 @@ struct ConvertVsstbSubviewOperandPattern adaptor.getMask()}); FailureOr> resultTypes = convertTypes(*getTypeConverter(), op->getResultTypes()); - if (failed(resultTypes)) + if (failed(resultTypes)) { return failure(); + } state.addTypes(*resultTypes); state.addAttributes(op->getAttrs()); Operation *newOp = rewriter.create(state); @@ -527,11 +590,13 @@ struct ConvertLoadScalarOperandToPtrPattern matchAndRewrite(pto::LoadScalarOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { Value ptr = materializeScalarAccessPtr(adaptor.getPtr(), rewriter, op.getLoc()); - if (!ptr) + if (!ptr) { return rewriter.notifyMatchFailure(op, "failed to materialize load_scalar ptr"); - if (!isa(ptr.getType())) + } + if (!isa(ptr.getType())) { return rewriter.notifyMatchFailure(op, "expected ptr-form load_scalar input"); + } rewriter.replaceOpWithNewOp(op, op.getValue().getType(), ptr, adaptor.getOffset()); @@ -547,11 +612,13 @@ struct ConvertStoreScalarOperandToPtrPattern matchAndRewrite(pto::StoreScalarOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { Value ptr = materializeScalarAccessPtr(adaptor.getPtr(), rewriter, op.getLoc()); - if (!ptr) + if (!ptr) { return rewriter.notifyMatchFailure(op, "failed to materialize store_scalar ptr"); - if (!isa(ptr.getType())) + } + if (!isa(ptr.getType())) { return rewriter.notifyMatchFailure(op, "expected ptr-form store_scalar input"); + } rewriter.replaceOpWithNewOp(op, ptr, adaptor.getOffset(), @@ -741,10 +808,12 @@ struct ConvertLoadOperandToPtrPattern : public OpConversionPattern(ptr.getType())) + } + if (!isa(ptr.getType())) { return rewriter.notifyMatchFailure(op, "expected ptr-form load input"); + } rewriter.replaceOpWithNewOp( op, op.getValue().getType(), ptr, adaptor.getOffset()); @@ -760,10 +829,12 @@ struct ConvertStoreOperandToPtrPattern matchAndRewrite(pto::PTOStoreOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { Value ptr = materializeScalarAccessPtr(adaptor.getPtr(), rewriter, op.getLoc()); - if (!ptr) + if (!ptr) { return rewriter.notifyMatchFailure(op, "failed to materialize store ptr"); - if (!isa(ptr.getType())) + } + if (!isa(ptr.getType())) { return rewriter.notifyMatchFailure(op, "expected ptr-form store input"); + } rewriter.replaceOpWithNewOp( op, ptr, adaptor.getOffset(), adaptor.getValue()); @@ -789,9 +860,10 @@ struct ConvertSimtLaunchOp final : public OpConversionPattern // emission validation. if (Value normalized = materializeScalarAccessPtr(originalArg, rewriter, op.getLoc())) { - if (normalized.getType() != ptrType) + if (normalized.getType() != ptrType) { normalized = rewriter.create(op.getLoc(), ptrType, normalized); + } arg = normalized; } } @@ -816,27 +888,33 @@ struct ConvertPtrNormalizeUnrealizedCastOp final LogicalResult matchAndRewrite(UnrealizedConversionCastOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - if (op->getNumOperands() != 1 || op->getNumResults() != 1) + if (op->getNumOperands() != 1 || op->getNumResults() != 1) { return failure(); + } if (!hasPtrNormalizeConvertibleType(op->getOperandTypes()) && - !hasPtrNormalizeConvertibleType(op->getResultTypes())) + !hasPtrNormalizeConvertibleType(op->getResultTypes())) { return failure(); + } Type convertedResultType = getTypeConverter()->convertType(op.getResult(0).getType()); - if (!convertedResultType) + if (!convertedResultType) { return failure(); + } Value input = adaptor.getOperands().front(); if (input.getType() != convertedResultType) { auto ptrType = dyn_cast(convertedResultType); - if (!ptrType) + if (!ptrType) { return failure(); + } input = materializeScalarAccessPtr(input, rewriter, op.getLoc()); - if (!input) + if (!input) { return failure(); - if (input.getType() != ptrType) + } + if (input.getType() != ptrType) { input = rewriter.create(op.getLoc(), ptrType, input); + } } rewriter.replaceOp(op, input); @@ -852,12 +930,14 @@ struct ConvertPtrNormalizeMemRefCastOp final matchAndRewrite(memref::CastOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { if (!hasPtrNormalizeConvertibleType(op.getSource().getType()) && - !hasPtrNormalizeConvertibleType(op.getType())) + !hasPtrNormalizeConvertibleType(op.getType())) { return failure(); + } Type convertedResultType = getTypeConverter()->convertType(op.getType()); - if (!convertedResultType) + if (!convertedResultType) { return failure(); + } Value source = adaptor.getSource(); if (source.getType() == convertedResultType) { @@ -866,9 +946,10 @@ struct ConvertPtrNormalizeMemRefCastOp final } if (!isa(source.getType()) || - !isa(convertedResultType)) + !isa(convertedResultType)) { return rewriter.notifyMatchFailure( op, "expected ptr/int memref.cast after ptr normalization"); + } rewriter.replaceOpWithNewOp(op, convertedResultType, source); @@ -1049,8 +1130,9 @@ struct VPTOPtrNormalizePass ConvertPtrNormalizeMemRefCastOp>( typeConverter, context); - if (failed(applyPartialConversion(module, target, std::move(patterns)))) + if (failed(applyPartialConversion(module, target, std::move(patterns)))) { signalPassFailure(); + } } }; diff --git a/lib/PTO/Transforms/VPTOSoftPostUpdate.cpp b/lib/PTO/Transforms/VPTOSoftPostUpdate.cpp index 217ca28d20..17709f2bab 100644 --- a/lib/PTO/Transforms/VPTOSoftPostUpdate.cpp +++ b/lib/PTO/Transforms/VPTOSoftPostUpdate.cpp @@ -36,7 +36,6 @@ namespace pto { using namespace mlir; namespace { - // A hardware block is 32 bytes; block-strided ops count in these units. static constexpr int64_t kBlockSizeBytes = 32; @@ -116,18 +115,23 @@ static const PostUpdateTable &getPostUpdateTable() { // plain byte-sized int/float rather than guess. static std::optional addPtrUnitBytes(Value base) { Type elemTy; - if (auto ptrTy = dyn_cast(base.getType())) + if (auto ptrTy = dyn_cast(base.getType())) { elemTy = ptrTy.getElementType(); - else if (auto memrefTy = dyn_cast(base.getType())) + } + else if (auto memrefTy = dyn_cast(base.getType())) { elemTy = memrefTy.getElementType(); - else + } + else { return std::nullopt; + } - if (!elemTy || !elemTy.isIntOrFloat()) + if (!elemTy || !elemTy.isIntOrFloat()) { return std::nullopt; + } unsigned bits = elemTy.getIntOrFloatBitWidth(); - if (bits == 0 || bits % 8 != 0) + if (bits == 0 || bits % 8 != 0) { return std::nullopt; + } return static_cast(bits / 8); } @@ -150,8 +154,9 @@ static std::optional strideUnitBytes(Operation *op, StrideUnit unit, static const PostUpdateOpInfo *getPostUpdateInfo(Operation *op) { auto it = getPostUpdateTable().find(op->getName().getStringRef()); - if (it == getPostUpdateTable().end()) + if (it == getPostUpdateTable().end()) { return nullptr; + } return &it->second; } @@ -215,8 +220,9 @@ static StrideExprRef makeLeaf(Value v) { // Compile-time value of `e`, if it has one. static std::optional foldConst(const StrideExprRef &e) { - if (!e) + if (!e) { return std::nullopt; + } switch (e->kind) { case StrideExpr::Kind::Const: return e->constant; @@ -231,12 +237,15 @@ static std::optional foldConst(const StrideExprRef &e) { case StrideExpr::Kind::Mul: { auto a = foldConst(e->lhs); auto b = foldConst(e->rhs); - if (!a || !b) + if (!a || !b) { return std::nullopt; - if (e->kind == StrideExpr::Kind::Add) + } + if (e->kind == StrideExpr::Kind::Add) { return *a + *b; - if (e->kind == StrideExpr::Kind::Sub) + } + if (e->kind == StrideExpr::Kind::Sub) { return *a - *b; + } return *a * *b; } } @@ -254,40 +263,50 @@ static StrideExprRef makeBinary(StrideExpr::Kind kind, StrideExprRef a, static StrideExprRef makeAdd(StrideExprRef a, StrideExprRef b) { auto ca = foldConst(a), cb = foldConst(b); - if (ca && cb) + if (ca && cb) { return makeConst(*ca + *cb); - if (ca && *ca == 0) + } + if (ca && *ca == 0) { return b; - if (cb && *cb == 0) + } + if (cb && *cb == 0) { return a; + } return makeBinary(StrideExpr::Kind::Add, a, b); } static StrideExprRef makeSub(StrideExprRef a, StrideExprRef b) { auto ca = foldConst(a), cb = foldConst(b); - if (ca && cb) + if (ca && cb) { return makeConst(*ca - *cb); - if (cb && *cb == 0) + } + if (cb && *cb == 0) { return a; + } return makeBinary(StrideExpr::Kind::Sub, a, b); } static StrideExprRef makeMul(StrideExprRef a, StrideExprRef b) { auto ca = foldConst(a), cb = foldConst(b); - if (ca && cb) + if (ca && cb) { return makeConst(*ca * *cb); - if ((ca && *ca == 0) || (cb && *cb == 0)) + } + if ((ca && *ca == 0) || (cb && *cb == 0)) { return makeConst(0); - if (ca && *ca == 1) + } + if (ca && *ca == 1) { return b; - if (cb && *cb == 1) + } + if (cb && *cb == 1) { return a; + } return makeBinary(StrideExpr::Kind::Mul, a, b); } static StrideExprRef makeCast(Operation *castOp, StrideExprRef a) { - if (auto c = foldConst(a)) + if (auto c = foldConst(a)) { return makeConst(*c); + } auto e = std::make_shared(); e->kind = StrideExpr::Kind::Cast; e->castOp = castOp; @@ -309,12 +328,15 @@ struct AffineForm { }; static bool sameAffineAtom(const StrideExprRef &a, const StrideExprRef &b) { - if (!a || !b || a->kind != b->kind) + if (!a || !b || a->kind != b->kind) { return false; - if (a->kind == StrideExpr::Kind::Leaf) + } + if (a->kind == StrideExpr::Kind::Leaf) { return a->leaf == b->leaf; - if (a->kind != StrideExpr::Kind::Cast) + } + if (a->kind != StrideExpr::Kind::Cast) { return false; + } return a->castOp->getName() == b->castOp->getName() && a->castOp->getOperand(0).getType() == b->castOp->getOperand(0).getType() && @@ -325,25 +347,31 @@ static bool sameAffineAtom(const StrideExprRef &a, const StrideExprRef &b) { static bool addAffineConstant(AffineForm &form, int64_t value) { int64_t result; - if (llvm::AddOverflow(form.constant, value, result)) + if (llvm::AddOverflow(form.constant, value, result)) { return false; + } form.constant = result; return true; } static bool addAffineTerm(AffineForm &form, StrideExprRef atom, int64_t coeff) { - if (coeff == 0) + if (coeff == 0) { return true; + } for (unsigned i = 0; i < form.terms.size(); ++i) { - if (!sameAffineAtom(form.terms[i].atom, atom)) + if (!sameAffineAtom(form.terms[i].atom, atom)) { continue; + } int64_t result; - if (llvm::AddOverflow(form.terms[i].coeff, coeff, result)) + if (llvm::AddOverflow(form.terms[i].coeff, coeff, result)) { return false; - if (result == 0) + } + if (result == 0) { form.terms.erase(form.terms.begin() + i); - else + } + else { form.terms[i].coeff = result; + } return true; } form.terms.push_back({std::move(atom), coeff}); @@ -352,8 +380,9 @@ static bool addAffineTerm(AffineForm &form, StrideExprRef atom, int64_t coeff) { static bool accumulateAffine(const StrideExprRef &e, int64_t scale, AffineForm &form) { - if (!e) + if (!e) { return false; + } switch (e->kind) { case StrideExpr::Kind::Const: { int64_t scaled; @@ -403,21 +432,24 @@ static bool accumulateAffine(const StrideExprRef &e, int64_t scale, static std::optional normalizeAffine(const StrideExprRef &e) { AffineForm form; - if (!accumulateAffine(e, 1, form)) + if (!accumulateAffine(e, 1, form)) { return std::nullopt; + } return form; } static bool equalAffineForms(const AffineForm &a, const AffineForm &b) { - if (a.constant != b.constant || a.terms.size() != b.terms.size()) + if (a.constant != b.constant || a.terms.size() != b.terms.size()) { return false; + } for (const AffineTerm &termA : a.terms) { - bool found = llvm::any_of(b.terms, [&](const AffineTerm &termB) { + bool found = llvm::any_of(b.terms, [&termA](const AffineTerm &termB) { return termA.coeff == termB.coeff && sameAffineAtom(termA.atom, termB.atom); }); - if (!found) + if (!found) { return false; + } } return true; } @@ -427,33 +459,39 @@ static bool isZeroAffineForm(const AffineForm &form) { } static bool divideAffineForm(AffineForm &form, int64_t divisor) { - if (divisor <= 0 || form.constant % divisor != 0) + if (divisor <= 0 || form.constant % divisor != 0) { return false; + } for (const AffineTerm &term : form.terms) - if (term.coeff % divisor != 0) + if (term.coeff % divisor != 0) { return false; + } form.constant /= divisor; - for (AffineTerm &term : form.terms) + for (AffineTerm &term : form.terms) { term.coeff /= divisor; + } return true; } static StrideExprRef affineFormToExpr(const AffineForm &form) { StrideExprRef result; - if (form.constant != 0) + if (form.constant != 0) { result = makeConst(form.constant); + } for (const AffineTerm &term : form.terms) { StrideExprRef value = term.atom; - if (term.coeff != 1) + if (term.coeff != 1) { value = makeMul(value, makeConst(term.coeff)); + } result = result ? makeAdd(result, value) : value; } return result ? result : makeConst(0); } static void collectLeaves(const StrideExprRef &e, SmallVectorImpl &out) { - if (!e) + if (!e) { return; + } if (e->kind == StrideExpr::Kind::Leaf) { out.push_back(e->leaf); return; @@ -479,10 +517,12 @@ static bool exprType(const StrideExprRef &e, Type &out) { case StrideExpr::Kind::Sub: case StrideExpr::Kind::Mul: { Type ta, tb; - if (!exprType(e->lhs, ta) || !exprType(e->rhs, tb)) + if (!exprType(e->lhs, ta) || !exprType(e->rhs, tb)) { return false; - if (ta && tb && ta != tb) + } + if (ta && tb && ta != tb) { return false; + } out = ta ? ta : tb; return true; } @@ -508,21 +548,24 @@ static std::optional decomposeLinear(Value v, scf::ForOp forOp, DecompCache &cache) { // v == blockArg → {1, 0} - if (v == blockArg) + if (v == blockArg) { return LinearDecomp{1, makeConst(0)}; + } auto it = cache.find(v); - if (it != cache.end()) + if (it != cache.end()) { return it->second; - auto record = [&](std::optional r) { + } + auto record = [&cache, &v](std::optional r) { cache[v] = r; return r; }; // v is other block arg (IV, different iter_arg, func arg) → {0, v} Operation *defOp = v.getDefiningOp(); - if (!defOp) + if (!defOp) { return record(LinearDecomp{0, makeLeaf(v)}); + } // v is loop-invariant or constant → {0, v} if (forOp.isDefinedOutsideOfLoop(v) || @@ -534,10 +577,12 @@ static std::optional decomposeLinear(Value v, if (isa(defOp)) { auto da = decomposeLinear(defOp->getOperand(0), blockArg, forOp, cache); auto db = decomposeLinear(defOp->getOperand(1), blockArg, forOp, cache); - if (!da || !db) + if (!da || !db) { return record(std::nullopt); - if (da->coeff == 0 && db->coeff == 0) + } + if (da->coeff == 0 && db->coeff == 0) { return record(LinearDecomp{0, makeLeaf(v)}); + } bool isSub = isa(defOp); return record( LinearDecomp{isSub ? da->coeff - db->coeff : da->coeff + db->coeff, @@ -549,17 +594,21 @@ static std::optional decomposeLinear(Value v, if (auto mulOp = dyn_cast(defOp)) { auto da = decomposeLinear(mulOp.getLhs(), blockArg, forOp, cache); auto db = decomposeLinear(mulOp.getRhs(), blockArg, forOp, cache); - if (!da || !db) + if (!da || !db) { return record(std::nullopt); - if (da->coeff == 0 && db->coeff == 0) + } + if (da->coeff == 0 && db->coeff == 0) { return record(LinearDecomp{0, makeLeaf(v)}); - if (da->coeff != 0 && db->coeff != 0) + } + if (da->coeff != 0 && db->coeff != 0) { return record(std::nullopt); + } const LinearDecomp &withBA = (da->coeff != 0) ? *da : *db; Value multiplier = (da->coeff != 0) ? mulOp.getRhs() : mulOp.getLhs(); auto constMul = getConstantIntValue(multiplier); - if (!constMul) + if (!constMul) { return record(std::nullopt); + } return record( LinearDecomp{withBA.coeff * *constMul, makeMul(withBA.increment, makeConst(*constMul))}); @@ -568,22 +617,27 @@ static std::optional decomposeLinear(Value v, // v = index_cast(a) → {ca, cast(ia)} when the cast preserves loop delta if (isa(defOp)) { auto d = decomposeLinear(defOp->getOperand(0), blockArg, forOp, cache); - if (!d) + if (!d) { return record(std::nullopt); - if (d->coeff == 0) + } + if (d->coeff == 0) { return record(LinearDecomp{0, makeLeaf(v)}); - if (!castPreservesLoopDelta(defOp, forOp)) + } + if (!castPreservesLoopDelta(defOp, forOp)) { return record(std::nullopt); + } return record(LinearDecomp{d->coeff, makeCast(defOp, d->increment)}); } // v = addptr(ptr, offset) → {c_ptr, i_ptr + offset} if (auto addPtrOp = dyn_cast(defOp)) { auto dp = decomposeLinear(addPtrOp.getPtr(), blockArg, forOp, cache); - if (!dp) + if (!dp) { return record(std::nullopt); - if (dp->coeff == 0) + } + if (dp->coeff == 0) { return record(LinearDecomp{0, makeLeaf(v)}); + } return record(LinearDecomp{ dp->coeff, makeAdd(dp->increment, makeLeaf(addPtrOp.getOffset()))}); } @@ -624,12 +678,14 @@ static StrideResult getIterArgIncrement(Value v, scf::ForOp forOp) { DecompCache cache; auto decomp = decomposeLinear(yieldOp.getOperand(idx), blockArg, forOp, cache); - if (!decomp || decomp->coeff != 1) + if (!decomp || decomp->coeff != 1) { return {StrideStatus::Failed, nullptr}; + } StrideExprRef inc = decomp->increment; - for (Operation *op : llvm::reverse(casts)) + for (Operation *op : llvm::reverse(casts)) { inc = makeCast(op, inc); + } return {StrideStatus::Ok, inc}; } @@ -639,8 +695,9 @@ static StrideResult getIterArgIncrement(Value v, scf::ForOp forOp) { return {StrideStatus::NotIterArg, nullptr}; if (isa(defOp)) { - if (!castPreservesLoopDelta(defOp, forOp)) + if (!castPreservesLoopDelta(defOp, forOp)) { return {StrideStatus::Failed, nullptr}; + } casts.push_back(defOp); current = defOp->getOperand(0); continue; @@ -684,8 +741,9 @@ static StrideResult getIterArgIncrement(Value v, scf::ForOp forOp) { using DeltaCache = DenseMap; static unsigned integerLikeBitWidth(Type type) { - if (type.isIndex()) + if (type.isIndex()) { return 64; + } return cast(type).getWidth(); } @@ -693,22 +751,26 @@ static std::optional getConstantTripCount(scf::ForOp forOp) { auto lower = getConstantIntValue(forOp.getLowerBound()); auto upper = getConstantIntValue(forOp.getUpperBound()); auto step = getConstantIntValue(forOp.getStep()); - if (!lower || !upper || !step || *step <= 0) + if (!lower || !upper || !step || *step <= 0) { return std::nullopt; - if (*lower >= *upper) + } + if (*lower >= *upper) { return 0; + } __int128 distance = static_cast<__int128>(*upper) - *lower; __int128 count = (distance + static_cast<__int128>(*step) - 1) / *step; - if (count > std::numeric_limits::max()) + if (count > std::numeric_limits::max()) { return std::nullopt; + } return static_cast(count); } static std::optional getConstantAPInt(Value value) { APInt constant; - if (!matchPattern(value, m_ConstantInt(&constant))) + if (!matchPattern(value, m_ConstantInt(&constant))) { return std::nullopt; + } return constant; } @@ -724,15 +786,18 @@ getConstantIterArgIncrement(BlockArgument iterArg, scf::ForOp forOp, unsigned idx = iterArg.getArgNumber() - 1; auto yieldOp = cast(forOp.getBody()->getTerminator()); Value yielded = yieldOp.getOperand(idx); - if (yielded == iterArg) + if (yielded == iterArg) { return 0; + } if (auto addOp = yielded.getDefiningOp()) { Value increment; - if (addOp.getLhs() == iterArg) + if (addOp.getLhs() == iterArg) { increment = addOp.getRhs(); - else if (addOp.getRhs() == iterArg) + } + else if (addOp.getRhs() == iterArg) { increment = addOp.getLhs(); + } auto constant = increment ? getConstantAPInt(increment) : std::nullopt; if (constant) return isUnsigned ? static_cast<__int128>(constant->getZExtValue()) @@ -740,8 +805,9 @@ getConstantIterArgIncrement(BlockArgument iterArg, scf::ForOp forOp, } if (auto subOp = yielded.getDefiningOp()) { - if (isUnsigned || subOp.getLhs() != iterArg) + if (isUnsigned || subOp.getLhs() != iterArg) { return std::nullopt; + } auto constant = getConstantAPInt(subOp.getRhs()); if (constant) { int64_t signedConstant = constant->getSExtValue(); @@ -763,14 +829,17 @@ static bool canonicalAddressRecurrenceDoesNotWrap(Value input, Operation *castOp, scf::ForOp forOp) { auto inputType = dyn_cast(input.getType()); - if (!inputType || inputType.getWidth() != kCanonicalAddressWidth) + if (!inputType || inputType.getWidth() != kCanonicalAddressWidth) { return false; + } auto tripCount = getConstantTripCount(forOp); - if (!tripCount) + if (!tripCount) { return false; - if (*tripCount == 0) + } + if (*tripCount == 0) { return true; + } bool isUnsigned = isa(castOp); APInt initialBits; @@ -778,8 +847,9 @@ static bool canonicalAddressRecurrenceDoesNotWrap(Value input, if (input == forOp.getInductionVar()) { auto lower = getConstantAPInt(forOp.getLowerBound()); auto step = getConstantAPInt(forOp.getStep()); - if (!lower || !step) + if (!lower || !step) { return false; + } initialBits = *lower; increment = static_cast<__int128>(step->getSExtValue()); } else { @@ -790,8 +860,9 @@ static bool canonicalAddressRecurrenceDoesNotWrap(Value input, unsigned idx = iterArg.getArgNumber() - 1; auto initial = getConstantAPInt(forOp.getInitArgs()[idx]); auto step = getConstantIterArgIncrement(iterArg, forOp, isUnsigned); - if (!initial || !step) + if (!initial || !step) { return false; + } initialBits = *initial; increment = *step; } @@ -828,17 +899,20 @@ static bool castPreservesLoopDelta(Operation *castOp, scf::ForOp forOp) { unsigned inputWidth = integerLikeBitWidth(input.getType()); unsigned resultWidth = integerLikeBitWidth(castOp->getResult(0).getType()); - if (resultWidth >= inputWidth) + if (resultWidth >= inputWidth) { return true; + } - if (input != forOp.getInductionVar()) + if (input != forOp.getInductionVar()) { return false; + } auto lower = getConstantIntValue(forOp.getLowerBound()); auto upper = getConstantIntValue(forOp.getUpperBound()); auto step = getConstantIntValue(forOp.getStep()); - if (!lower || !upper || !step || *step <= 0) + if (!lower || !upper || !step || *step <= 0) { return false; + } if (*lower >= *upper) return true; // Empty loop. @@ -855,17 +929,20 @@ static bool castPreservesLoopDelta(Operation *castOp, scf::ForOp forOp) { static StrideExprRef computeDelta(Value v, scf::ForOp forOp, DeltaCache &cache) { // IV: delta = step - if (v == forOp.getInductionVar()) + if (v == forOp.getInductionVar()) { return makeLeaf(forOp.getStep()); + } // Constant or loop-invariant: delta = 0 - if (forOp.isDefinedOutsideOfLoop(v)) + if (forOp.isDefinedOutsideOfLoop(v)) { return makeConst(0); + } auto it = cache.find(v); - if (it != cache.end()) + if (it != cache.end()) { return it->second; - auto record = [&](StrideExprRef r) { + } + auto record = [&cache, &v](StrideExprRef r) { cache[v] = r; return r; }; @@ -878,27 +955,32 @@ static StrideExprRef computeDelta(Value v, scf::ForOp forOp, Value yieldVal = yieldOp.getOperand(idx); if (auto addOp = yieldVal.getDefiningOp()) { Value other; - if (addOp.getLhs() == blockArg) + if (addOp.getLhs() == blockArg) { other = addOp.getRhs(); - else if (addOp.getRhs() == blockArg) + } + else if (addOp.getRhs() == blockArg) { other = addOp.getLhs(); - if (other && forOp.isDefinedOutsideOfLoop(other)) + } + if (other && forOp.isDefinedOutsideOfLoop(other)) { return record(makeLeaf(other)); + } } return record(nullptr); } } Operation *defOp = v.getDefiningOp(); - if (!defOp) + if (!defOp) { return record(nullptr); + } // arith.addi(a, b): delta = delta(a) + delta(b) if (auto addOp = dyn_cast(defOp)) { auto da = computeDelta(addOp.getLhs(), forOp, cache); auto db = computeDelta(addOp.getRhs(), forOp, cache); - if (!da || !db) + if (!da || !db) { return record(nullptr); + } return record(makeAdd(da, db)); } @@ -906,8 +988,9 @@ static StrideExprRef computeDelta(Value v, scf::ForOp forOp, if (auto subOp = dyn_cast(defOp)) { auto da = computeDelta(subOp.getLhs(), forOp, cache); auto db = computeDelta(subOp.getRhs(), forOp, cache); - if (!da || !db) + if (!da || !db) { return record(nullptr); + } return record(makeSub(da, db)); } @@ -919,8 +1002,9 @@ static StrideExprRef computeDelta(Value v, scf::ForOp forOp, {std::pair{rhs, lhs}, std::pair{lhs, rhs}}) { if (forOp.isDefinedOutsideOfLoop(invariant)) { auto dv = computeDelta(variant, forOp, cache); - if (!dv) + if (!dv) { continue; + } return record(makeMul(makeLeaf(invariant), dv)); } } @@ -930,8 +1014,9 @@ static StrideExprRef computeDelta(Value v, scf::ForOp forOp, // Preserve value-preserving casts in the symbolic delta. Narrowing casts are // accepted only when castPreservesLoopDelta proves they cannot truncate. if (isa(defOp)) { - if (!castPreservesLoopDelta(defOp, forOp)) + if (!castPreservesLoopDelta(defOp, forOp)) { return record(nullptr); + } StrideExprRef inputDelta = computeDelta(defOp->getOperand(0), forOp, cache); return record(inputDelta ? makeCast(defOp, inputDelta) : nullptr); } @@ -945,10 +1030,12 @@ static StrideExprRef computeDelta(Value v, scf::ForOp forOp, // Returns null if the stride cannot be determined. static StrideExprRef getStride(Value v, scf::ForOp forOp, DeltaCache &cache) { StrideResult r = getIterArgIncrement(v, forOp); - if (r.status == StrideStatus::Ok) + if (r.status == StrideStatus::Ok) { return r.expr; - if (r.status == StrideStatus::Failed) + } + if (r.status == StrideStatus::Failed) { return nullptr; + } return computeDelta(v, forOp, cache); } @@ -962,12 +1049,14 @@ static StrideExprRef getStride(Value v, scf::ForOp forOp, DeltaCache &cache) { static Value materializeAtLoopEntry(Value v, scf::ForOp forOp, OpBuilder &builder) { // IV → lower bound - if (v == forOp.getInductionVar()) + if (v == forOp.getInductionVar()) { return forOp.getLowerBound(); + } // Already defined outside the loop — use directly. - if (forOp.isDefinedOutsideOfLoop(v)) + if (forOp.isDefinedOutsideOfLoop(v)) { return v; + } // iter_arg → its init value if (auto blockArg = dyn_cast(v)) { @@ -978,26 +1067,30 @@ static Value materializeAtLoopEntry(Value v, scf::ForOp forOp, } Operation *defOp = v.getDefiningOp(); - if (!defOp || !forOp->isAncestor(defOp)) + if (!defOp || !forOp->isAncestor(defOp)) { return nullptr; + } // Cloning duplicates the op, so it must be safe to execute an extra time and // its result must not depend on anything but its operands. - if (!isPure(defOp)) + if (!isPure(defOp)) { return nullptr; + } // Clone the defining op with operands materialized at loop entry. SmallVector newOperands; for (Value operand : defOp->getOperands()) { Value materialized = materializeAtLoopEntry(operand, forOp, builder); - if (!materialized) + if (!materialized) { return nullptr; + } newOperands.push_back(materialized); } builder.setInsertionPoint(forOp); Operation *cloned = builder.clone(*defOp); - for (auto [i, operand] : llvm::enumerate(newOperands)) + for (auto [i, operand] : llvm::enumerate(newOperands)) { cloned->setOperand(i, operand); + } // Preserve which result was asked for; `v` need not be result 0. return cloned->getResult(cast(v).getResultNumber()); } @@ -1007,10 +1100,12 @@ static Value materializeAtLoopEntry(Value v, scf::ForOp forOp, // types; finer byte units require a suitably aligned compile-time constant. static bool canScaleInitialOffset(Value strideOperand, int64_t elemBytes, int64_t unitBytes) { - if (!strideOperand || unitBytes == elemBytes || unitBytes % elemBytes == 0) + if (!strideOperand || unitBytes == elemBytes || unitBytes % elemBytes == 0) { return true; - if (elemBytes % unitBytes != 0) + } + if (elemBytes % unitBytes != 0) { return false; + } auto constant = getConstantIntValue(strideOperand); return constant && *constant % (elemBytes / unitBytes) == 0; } @@ -1022,8 +1117,9 @@ static bool canScaleInitialOffset(Value strideOperand, int64_t elemBytes, // offsets in the same way as the existing lowering. static Value truncateElementOffsetToI32(Value offset, Location loc, OpBuilder &builder) { - if (!offset.getType().isIndex()) + if (!offset.getType().isIndex()) { return offset; + } Value offsetI64 = builder.create(loc, builder.getI64Type(), offset); Value offsetI32 = @@ -1037,8 +1133,9 @@ static Value truncateElementOffsetToI32(Value offset, Location loc, // signed, including sprsti's signed 8-bit word offset. static Value normalizeAddPtrOffsetToIndex(Value offset, StrideUnit strideUnit, Location loc, OpBuilder &builder) { - if (offset.getType().isIndex()) + if (offset.getType().isIndex()) { return offset; + } if (strideUnit == StrideUnit::Block) return builder.create(loc, builder.getIndexType(), offset); @@ -1052,13 +1149,16 @@ static Value createInitialPtr(Value base, Value strideOperand, StrideUnit strideUnit, int64_t elemBytes, int64_t unitBytes, Location loc, OpBuilder &builder) { - if (!strideOperand) + if (!strideOperand) { return base; + } auto constSo = getConstantIntValue(strideOperand); - if (constSo && *constSo == 0) + if (constSo && *constSo == 0) { return base; - if (!canScaleInitialOffset(strideOperand, elemBytes, unitBytes)) + } + if (!canScaleInitialOffset(strideOperand, elemBytes, unitBytes)) { return nullptr; + } Value scaledOffset = strideUnit == StrideUnit::Element @@ -1092,15 +1192,18 @@ static Value computeInitialPtr(Value base, Value strideOperand, int64_t unitBytes, scf::ForOp forOp, OpBuilder &builder) { Value baseAtEntry = materializeAtLoopEntry(base, forOp, builder); - if (!baseAtEntry) + if (!baseAtEntry) { return nullptr; + } - if (!strideOperand) + if (!strideOperand) { return baseAtEntry; + } Value soAtEntry = materializeAtLoopEntry(strideOperand, forOp, builder); - if (!soAtEntry) + if (!soAtEntry) { return nullptr; + } builder.setInsertionPoint(forOp); return createInitialPtr(baseAtEntry, soAtEntry, strideUnit, elemBytes, @@ -1109,7 +1212,6 @@ static Value computeInitialPtr(Value base, Value strideOperand, // Rescale a per-iteration base delta from `pto.addptr` units (elements) into // the op's strideOperand units. Returns null when the conversion is not exact. -// // Expanding the byte-denominated form // stride_new = (E*delta(base) + W*delta(strideOperand)) / W // = (E/W)*delta(base) + delta(strideOperand) @@ -1120,21 +1222,24 @@ static Value computeInitialPtr(Value base, Value strideOperand, // behave exactly as before this scaling existed. static StrideExprRef scaleBaseDelta(StrideExprRef deltaBase, int64_t elemBytes, int64_t unitBytes) { - if (unitBytes == elemBytes) + if (unitBytes == elemBytes) { return deltaBase; + } if (unitBytes % elemBytes == 0) { // Coarser stride unit (e.g. 32-byte blocks over 4-byte elements): the base // delta's affine coefficients must all land on a whole unit. int64_t divisor = unitBytes / elemBytes; auto form = normalizeAffine(deltaBase); - if (!form || !divideAffineForm(*form, divisor)) + if (!form || !divideAffineForm(*form, divisor)) { return nullptr; + } return affineFormToExpr(*form); } - if (elemBytes % unitBytes == 0) + if (elemBytes % unitBytes == 0) { return makeMul(deltaBase, makeConst(elemBytes / unitBytes)); + } return nullptr; } @@ -1149,17 +1254,20 @@ static StrideExprRef combineStride(StrideExprRef deltaBase, StrideExprRef deltaOffset, int64_t elemBytes, int64_t unitBytes) { StrideExprRef scaledBase = scaleBaseDelta(deltaBase, elemBytes, unitBytes); - if (!scaledBase) + if (!scaledBase) { return nullptr; + } StrideExprRef total = makeAdd(scaledBase, deltaOffset); if (auto form = normalizeAffine(total)) { - if (isZeroAffineForm(*form)) + if (isZeroAffineForm(*form)) { return nullptr; + } return affineFormToExpr(*form); } - if (auto constTotal = foldConst(total); constTotal && *constTotal == 0) + if (auto constTotal = foldConst(total); constTotal && *constTotal == 0) { return nullptr; + } return total; } @@ -1171,27 +1279,34 @@ static StrideExprRef combineStride(StrideExprRef deltaBase, // def-chain if needed? Pure: inspects only, never mutates the IR. static bool canHoistBefore(Value v, Operation *insertPt, scf::ForOp forOp, DenseMap &memo) { - if (forOp.isDefinedOutsideOfLoop(v) || isa(v)) + if (forOp.isDefinedOutsideOfLoop(v) || isa(v)) { return true; + } Operation *defOp = v.getDefiningOp(); - if (!defOp) + if (!defOp) { return true; - if (defOp->getBlock() != insertPt->getBlock()) + } + if (defOp->getBlock() != insertPt->getBlock()) { return false; + } // Already earlier in the block, so usable as-is: SSA guarantees its operands // are defined even earlier. This reasoning is only sound because analysis // creates no IR — every value examined here predates the transform. - if (defOp->isBeforeInBlock(insertPt)) + if (defOp->isBeforeInBlock(insertPt)) { return true; + } auto it = memo.find(v); - if (it != memo.end()) + if (it != memo.end()) { return it->second; + } memo[v] = false; - if (!isPure(defOp)) + if (!isPure(defOp)) { return false; + } for (Value operand : defOp->getOperands()) - if (!canHoistBefore(operand, insertPt, forOp, memo)) + if (!canHoistBefore(operand, insertPt, forOp, memo)) { return false; + } memo[v] = true; return true; } @@ -1200,24 +1315,28 @@ static bool canHoistBefore(Value v, Operation *insertPt, scf::ForOp forOp, // canHoistBefore has approved `v`. static Value hoistBefore(Value v, Operation *insertPt, scf::ForOp forOp, OpBuilder &builder, DenseMap &memo) { - if (forOp.isDefinedOutsideOfLoop(v) || isa(v)) + if (forOp.isDefinedOutsideOfLoop(v) || isa(v)) { return v; + } Operation *defOp = v.getDefiningOp(); if (!defOp || defOp->getBlock() != insertPt->getBlock() || defOp->isBeforeInBlock(insertPt)) return v; auto it = memo.find(v); - if (it != memo.end()) + if (it != memo.end()) { return it->second; + } SmallVector newOperands; - for (Value operand : defOp->getOperands()) + for (Value operand : defOp->getOperands()) { newOperands.push_back(hoistBefore(operand, insertPt, forOp, builder, memo)); + } builder.setInsertionPoint(insertPt); Operation *cloned = builder.clone(*defOp); - for (auto [i, operand] : llvm::enumerate(newOperands)) + for (auto [i, operand] : llvm::enumerate(newOperands)) { cloned->setOperand(i, operand); + } // Preserve which result was asked for; `v` need not be result 0. Value res = cloned->getResult(cast(v).getResultNumber()); memo[v] = res; @@ -1262,17 +1381,20 @@ using ConstCache = DenseMap, Value>; static Value materializeConst(int64_t c, Type ty, Location loc, scf::ForOp forOp, ConstCache &cache, OpBuilder &builder) { - if (!ty) + if (!ty) { ty = builder.getIndexType(); + } auto key = std::make_pair(c, ty); - if (auto it = cache.find(key); it != cache.end()) + if (auto it = cache.find(key); it != cache.end()) { return it->second; + } OpBuilder::InsertionGuard guard(builder); builder.setInsertionPoint(forOp); Value v; - if (ty.isIndex()) + if (ty.isIndex()) { v = builder.create(loc, c); + } else v = builder.create(loc, c, ty.getIntOrFloatBitWidth()); @@ -1308,11 +1430,13 @@ static bool constantsFitType(const StrideExprRef &e, Type wantType) { static bool satisfiesStrideConstraint(const StrideExprRef &stride, StrideConstraint constraint) { - if (constraint == StrideConstraint::Dynamic) + if (constraint == StrideConstraint::Dynamic) { return true; + } std::optional constant = foldConst(stride); - if (!constant) + if (!constant) { return false; + } return constraint == StrideConstraint::Constant || (*constant >= -128 && *constant <= 127); } @@ -1340,10 +1464,12 @@ static Value materialize(const StrideExprRef &e, Type wantType, Location loc, case StrideExpr::Kind::Mul: { Value a = materialize(e->lhs, wantType, loc, forOp, cache, builder); Value b = materialize(e->rhs, a.getType(), loc, forOp, cache, builder); - if (e->kind == StrideExpr::Kind::Add) + if (e->kind == StrideExpr::Kind::Add) { return builder.create(loc, a, b); - if (e->kind == StrideExpr::Kind::Sub) + } + if (e->kind == StrideExpr::Kind::Sub) { return builder.create(loc, a, b); + } return builder.create(loc, a, b); } } @@ -1382,29 +1508,35 @@ static bool canMaterializeAs(const StrideExprRef &e, Type wantType) { static bool canHoistBefore(Value v, Operation *insertPt, DominanceInfo &dominance, DenseMap &memo) { - if (dominance.dominates(v, insertPt)) + if (dominance.dominates(v, insertPt)) { return true; + } auto it = memo.find(v); - if (it != memo.end()) + if (it != memo.end()) { return it->second; + } memo[v] = false; Operation *defOp = v.getDefiningOp(); - if (!defOp || defOp->getBlock() != insertPt->getBlock() || !isPure(defOp)) + if (!defOp || defOp->getBlock() != insertPt->getBlock() || !isPure(defOp)) { return false; + } for (Value operand : defOp->getOperands()) - if (!canHoistBefore(operand, insertPt, dominance, memo)) + if (!canHoistBefore(operand, insertPt, dominance, memo)) { return false; + } memo[v] = true; return true; } static Value hoistBefore(Value v, Operation *insertPt, DominanceInfo &dominance, OpBuilder &builder, DenseMap &memo) { - if (dominance.dominates(v, insertPt)) + if (dominance.dominates(v, insertPt)) { return v; - if (auto it = memo.find(v); it != memo.end()) + } + if (auto it = memo.find(v); it != memo.end()) { return it->second; + } Operation *defOp = v.getDefiningOp(); SmallVector newOperands; @@ -1414,8 +1546,9 @@ static Value hoistBefore(Value v, Operation *insertPt, DominanceInfo &dominance, builder.setInsertionPoint(insertPt); Operation *cloned = builder.clone(*defOp); - for (auto [i, operand] : llvm::enumerate(newOperands)) + for (auto [i, operand] : llvm::enumerate(newOperands)) { cloned->setOperand(i, operand); + } Value result = cloned->getResult(cast(v).getResultNumber()); memo[v] = result; return result; @@ -1453,8 +1586,9 @@ static Value materializeSequential(const StrideExprRef &e, Type wantType, Location loc, OpBuilder &builder) { switch (e->kind) { case StrideExpr::Kind::Const: - if (wantType.isIndex()) + if (wantType.isIndex()) { return builder.create(loc, e->constant); + } return builder.create( loc, e->constant, wantType.getIntOrFloatBitWidth()); case StrideExpr::Kind::Leaf: @@ -1471,10 +1605,12 @@ static Value materializeSequential(const StrideExprRef &e, Type wantType, case StrideExpr::Kind::Mul: { Value lhs = materializeSequential(e->lhs, wantType, loc, builder); Value rhs = materializeSequential(e->rhs, wantType, loc, builder); - if (e->kind == StrideExpr::Kind::Add) + if (e->kind == StrideExpr::Kind::Add) { return builder.create(loc, lhs, rhs); - if (e->kind == StrideExpr::Kind::Sub) + } + if (e->kind == StrideExpr::Kind::Sub) { return builder.create(loc, lhs, rhs); + } return builder.create(loc, lhs, rhs); } } @@ -1520,15 +1656,19 @@ static Operation *createPostUpdateOp(Operation *op, Value stride, OpBuilder &builder) { OperationState state(op->getLoc(), op->getName()); for (auto [i, operand] : llvm::enumerate(op->getOperands())) { - if (static_cast(i) == info.baseOperandIdx) + if (static_cast(i) == info.baseOperandIdx) { state.addOperands(base); - else if (info.strideOperandIdx && i == *info.strideOperandIdx) + } + else if (info.strideOperandIdx && i == *info.strideOperandIdx) { state.addOperands(stride); - else + } + else { state.addOperands(operand); + } } - if (!info.strideOperandIdx) + if (!info.strideOperandIdx) { state.addOperands(stride); + } state.addTypes(op->getResultTypes()); state.addTypes(base.getType()); state.addAttributes(op->getAttrs()); @@ -1542,12 +1682,15 @@ static Operation *createNormalOp(Operation *op, const PostUpdateOpInfo &info, OpBuilder &builder) { OperationState state(op->getLoc(), op->getName()); for (auto [i, operand] : llvm::enumerate(op->getOperands())) { - if (static_cast(i) == info.baseOperandIdx) + if (static_cast(i) == info.baseOperandIdx) { state.addOperands(base); - else if (info.strideOperandIdx && i == *info.strideOperandIdx) + } + else if (info.strideOperandIdx && i == *info.strideOperandIdx) { state.addOperands(zeroStride); - else + } + else { state.addOperands(operand); + } } state.addTypes(op->getResultTypes()); state.addAttributes(op->getAttrs()); @@ -1555,12 +1698,9 @@ static Operation *createNormalOp(Operation *op, const PostUpdateOpInfo &info, } // Remove loop-carried recurrences that became dead after post-update rewrites. -// // Ordinary DCE cannot break a recurrence such as -// // %next = arith.addi %iter_arg, %step // scf.yield %next -// // even when neither the iter_arg nor the loop result has a real user: the // block argument, update, and yield form a use cycle. Compute liveness from // side-effecting operations and externally used loop results, close it across @@ -1569,23 +1709,26 @@ static Operation *createNormalOp(Operation *op, const PostUpdateOpInfo &info, static scf::ForOp pruneDeadLoopCarriedValues(scf::ForOp forOp, OpBuilder &builder) { unsigned numIterArgs = forOp.getInitArgs().size(); - if (numIterArgs == 0) + if (numIterArgs == 0) { return forOp; + } auto yieldOp = cast(forOp.getBody()->getTerminator()); DenseSet liveValues; SmallVector worklist; SmallVector keepIterArg(numIterArgs, false); - auto markLive = [&](Value value) { - if (value && liveValues.insert(value).second) + auto markLive = [&liveValues, &worklist](Value value) { + if (value && liveValues.insert(value).second) { worklist.push_back(value); + } }; // A loop result used outside the loop keeps its corresponding backedge live. for (auto [idx, result] : llvm::enumerate(forOp.getResults())) { - if (result.use_empty()) + if (result.use_empty()) { continue; + } keepIterArg[idx] = true; markLive(yieldOp.getOperand(idx)); } @@ -1597,12 +1740,14 @@ static scf::ForOp pruneDeadLoopCarriedValues(scf::ForOp forOp, // here. for (Operation &op : forOp.getBody()->without_terminator()) { if (!isPure(&op)) - for (Value operand : op.getOperands()) + for (Value operand : op.getOperands()) { markLive(operand); + } if (op.getNumRegions() != 0) { - op.walk([&](Operation *nested) { - for (Value operand : nested->getOperands()) + op.walk([&markLive](Operation *nested) { + for (Value operand : nested->getOperands()) { markLive(operand); + } }); } } @@ -1623,20 +1768,24 @@ static scf::ForOp pruneDeadLoopCarriedValues(scf::ForOp forOp, } Operation *defOp = value.getDefiningOp(); - if (!defOp || !forOp->isAncestor(defOp)) + if (!defOp || !forOp->isAncestor(defOp)) { continue; - for (Value operand : defOp->getOperands()) + } + for (Value operand : defOp->getOperands()) { markLive(operand); + } } - if (llvm::all_of(keepIterArg, [](bool keep) { return keep; })) + if (llvm::all_of(keepIterArg, [](bool keep) { return keep; })) { return forOp; + } SmallVector newInitArgs; newInitArgs.reserve(numIterArgs); for (auto [idx, init] : llvm::enumerate(forOp.getInitArgs())) - if (keepIterArg[idx]) + if (keepIterArg[idx]) { newInitArgs.push_back(init); + } builder.setInsertionPoint(forOp); auto newForOp = builder.create( @@ -1648,23 +1797,26 @@ static scf::ForOp pruneDeadLoopCarriedValues(scf::ForOp forOp, mapping.map(forOp.getInductionVar(), newForOp.getInductionVar()); unsigned newArgIdx = 0; for (auto [idx, oldArg] : llvm::enumerate(forOp.getRegionIterArgs())) - if (keepIterArg[idx]) + if (keepIterArg[idx]) { mapping.map(oldArg, newForOp.getRegionIterArgs()[newArgIdx++]); + } builder.setInsertionPointToStart(newForOp.getBody()); for (Operation &op : forOp.getBody()->without_terminator()) { - bool hasLiveResult = llvm::any_of(op.getResults(), [&](Value result) { + bool hasLiveResult = llvm::any_of(op.getResults(), [&liveValues](Value result) { return liveValues.contains(result); }); - if (!isPure(&op) || op.getNumRegions() != 0 || hasLiveResult) + if (!isPure(&op) || op.getNumRegions() != 0 || hasLiveResult) { builder.clone(op, mapping); + } } SmallVector newYields; newYields.reserve(newInitArgs.size()); for (auto [idx, yielded] : llvm::enumerate(yieldOp.getOperands())) - if (keepIterArg[idx]) + if (keepIterArg[idx]) { newYields.push_back(mapping.lookupOrDefault(yielded)); + } builder.setInsertionPointToEnd(newForOp.getBody()); builder.create(yieldOp.getLoc(), newYields); @@ -1686,8 +1838,9 @@ static scf::ForOp pruneDeadLoopCarriedValues(scf::ForOp forOp, static scf::ForOp applyPostUpdateRewrites(scf::ForOp forOp, ArrayRef rewrites, OpBuilder &builder) { - if (rewrites.empty()) + if (rewrites.empty()) { return nullptr; + } // Group rewrites by start-address operands, stride, and effective byte unit. // Ops in the same group share one iter_arg and all use the pre-update @@ -1702,8 +1855,9 @@ static scf::ForOp applyPostUpdateRewrites(scf::ForOp forOp, for (auto [i, rw] : llvm::enumerate(rewrites)) { auto key = getGroupKey(rw); auto [it, inserted] = groupToIdx.try_emplace(key, groupInitPtrs.size()); - if (inserted) + if (inserted) { groupInitPtrs.push_back(rw.initPtr); + } rwGroupIdx[i] = it->second; } @@ -1712,8 +1866,9 @@ static scf::ForOp applyPostUpdateRewrites(scf::ForOp forOp, // Build new init args: original + one new pointer per group. SmallVector newInitArgs(forOp.getInitArgs().begin(), forOp.getInitArgs().end()); - for (Value ptr : groupInitPtrs) + for (Value ptr : groupInitPtrs) { newInitArgs.push_back(ptr); + } unsigned origIterArgCount = forOp.getInitArgs().size(); @@ -1729,8 +1884,9 @@ static scf::ForOp applyPostUpdateRewrites(scf::ForOp forOp, Block *oldBody = forOp.getBody(); Block *newBody = newForOp.getBody(); mapping.map(forOp.getInductionVar(), newForOp.getInductionVar()); - for (unsigned i = 0; i < origIterArgCount; ++i) + for (unsigned i = 0; i < origIterArgCount; ++i) { mapping.map(oldBody->getArgument(i + 1), newBody->getArgument(i + 1)); + } // Clone the body, tracking old->new op correspondence. DenseMap opMapping; @@ -1743,13 +1899,15 @@ static scf::ForOp applyPostUpdateRewrites(scf::ForOp forOp, // Apply rewrites. All ops in a group use the same pre-update pointer (block // arg). Track the last updated_base per group for yielding. SmallVector groupYieldPtrs(numGroups); - for (unsigned g = 0; g < numGroups; ++g) + for (unsigned g = 0; g < numGroups; ++g) { groupYieldPtrs[g] = newBody->getArgument(origIterArgCount + 1 + g); + } for (auto [rwIdx, rw] : llvm::enumerate(rewrites)) { auto it = opMapping.find(rw.op); - if (it == opMapping.end()) + if (it == opMapping.end()) { continue; + } Operation *clonedOp = it->second; unsigned gIdx = rwGroupIdx[rwIdx]; Value ptr = newBody->getArgument(origIterArgCount + 1 + gIdx); @@ -1758,8 +1916,9 @@ static scf::ForOp applyPostUpdateRewrites(scf::ForOp forOp, builder.setInsertionPoint(clonedOp); const PostUpdateOpInfo *info = getPostUpdateInfo(clonedOp); - if (!info) + if (!info) { continue; + } Operation *newOp = createPostUpdateOp(clonedOp, *info, ptr, strideNew, builder); @@ -1780,17 +1939,20 @@ static scf::ForOp applyPostUpdateRewrites(scf::ForOp forOp, // Build yield: original yields + one pointer per group. auto oldYield = cast(oldBody->getTerminator()); SmallVector newYields; - for (Value v : oldYield.getOperands()) + for (Value v : oldYield.getOperands()) { newYields.push_back(mapping.lookupOrDefault(v)); - for (Value ptr : groupYieldPtrs) + } + for (Value ptr : groupYieldPtrs) { newYields.push_back(ptr); + } builder.setInsertionPointToEnd(newBody); builder.create(oldYield.getLoc(), newYields); // Replace original ForOp results (only the original ones). - for (unsigned i = 0; i < forOp.getNumResults(); ++i) + for (unsigned i = 0; i < forOp.getNumResults(); ++i) { forOp.getResult(i).replaceAllUsesWith(newForOp.getResult(i)); + } forOp.erase(); return pruneDeadLoopCarriedValues(newForOp, builder); @@ -1807,12 +1969,15 @@ using SequentialExprCache = DenseMap; // SSA value is reused without guessing at its semantics. static StrideExprRef buildSequentialExpr(Value value, SequentialExprCache &cache) { - if (!value) + if (!value) { return makeConst(0); - if (auto constant = getConstantIntValue(value)) + } + if (auto constant = getConstantIntValue(value)) { return makeConst(*constant); - if (auto it = cache.find(value); it != cache.end()) + } + if (auto it = cache.find(value); it != cache.end()) { return it->second; + } Operation *defOp = value.getDefiningOp(); StrideExprRef result; @@ -1829,8 +1994,9 @@ static StrideExprRef buildSequentialExpr(Value value, else if (auto rhsConst = getConstantIntValue(mul.getRhs())) result = makeMul(buildSequentialExpr(mul.getLhs(), cache), makeConst(*rhsConst)); - else + else { result = makeLeaf(value); + } } else if (isa_and_nonnull(defOp)) { result = makeCast(defOp, buildSequentialExpr(defOp->getOperand(0), cache)); } else { @@ -1852,8 +2018,9 @@ static NormalizedBase normalizeSequentialBase(Value base, int64_t elemBytes, StrideExprRef offset = makeConst(0); while (auto addPtr = root.getDefiningOp()) { auto parentElemBytes = addPtrUnitBytes(addPtr.getPtr()); - if (!parentElemBytes || *parentElemBytes != elemBytes) + if (!parentElemBytes || *parentElemBytes != elemBytes) { break; + } offset = makeAdd(offset, buildSequentialExpr(addPtr.getOffset(), cache)); root = addPtr.getPtr(); } @@ -1894,11 +2061,13 @@ analyzeSequentialStep(const SequentialCandidate &previous, StrideExprRef deltaStride = makeSub(current.strideExpr, previous.strideExpr); StrideExprRef step = combineStride(deltaBase, deltaStride, current.elemBytes, current.unitBytes); - if (!step) + if (!step) { return std::nullopt; + } auto form = normalizeAffine(step); - if (!form || isZeroAffineForm(*form)) + if (!form || isZeroAffineForm(*form)) { return std::nullopt; + } return SequentialStep{affineFormToExpr(*form), std::move(*form)}; } @@ -1914,8 +2083,9 @@ struct SequentialRun { static bool validateSequentialRun(SequentialRun &run, DominanceInfo &dominance) { - if (run.candidates.size() < 3) + if (run.candidates.size() < 3) { return false; + } SequentialCandidate *first = run.candidates.front(); run.strideType = first->strideOperand @@ -1937,8 +2107,9 @@ static bool validateSequentialRun(SequentialRun &run, candidate->strideOperand ? candidate->strideOperand.getType() : IndexType::get(candidate->op->getContext()); - if (candidateStrideType != run.strideType) + if (candidateStrideType != run.strideType) { return false; + } } SmallVector leaves; @@ -1969,10 +2140,12 @@ static unsigned countDeadDynamicAddPtrs(const SequentialRun &run) { Value value = candidate->base; Operation *expectedUser = candidate->op; while (auto addPtr = value.getDefiningOp()) { - if (!hasOnlyExpectedUser(value, expectedUser)) + if (!hasOnlyExpectedUser(value, expectedUser)) { break; - if (isDynamicSequentialValue(addPtr.getOffset(), cache)) + } + if (isDynamicSequentialValue(addPtr.getOffset(), cache)) { counted.insert(addPtr); + } expectedUser = addPtr; value = addPtr.getPtr(); } @@ -1982,14 +2155,15 @@ static unsigned countDeadDynamicAddPtrs(const SequentialRun &run) { static unsigned initialPointerCost(const SequentialRun &run) { SequentialCandidate *first = run.candidates.front(); - if (!first->info->strideIsInitialOffset || !first->strideOperand) + if (!first->info->strideIsInitialOffset || !first->strideOperand) { return 0; + } auto initialOffset = getConstantIntValue(first->strideOperand); return initialOffset && *initialOffset == 0 ? 0 : 1; } static bool isRunStrideUse(OpOperand &use, const SequentialRun &run) { - return llvm::any_of(run.candidates, [&](SequentialCandidate *candidate) { + return llvm::any_of(run.candidates, [&use](SequentialCandidate *candidate) { return candidate->info->strideOperandIdx && use.getOwner() == candidate->op && use.getOperandNumber() == @@ -2006,8 +2180,9 @@ static void collectCumulativeOffsetOps(Value value, if (!isa_and_nonnull(defOp) || !ops.insert(defOp).second) return; - for (Value operand : defOp->getOperands()) + for (Value operand : defOp->getOperands()) { collectCumulativeOffsetOps(operand, ops); + } } static bool allUsesDisappearAfterRewrite(Operation *op, @@ -2024,25 +2199,29 @@ static bool cumulativeOffsetChainDefinitelyDies( const SequentialRun &run, DenseSet &deadOps) { for (SequentialCandidate *candidate : llvm::drop_begin(run.candidates, 2)) { - if (candidate->strideOperand) + if (candidate->strideOperand) { collectCumulativeOffsetOps(candidate->strideOperand, deadOps); + } } return !deadOps.empty() && - llvm::all_of(deadOps, [&](Operation *op) { - return allUsesDisappearAfterRewrite(op, deadOps, run); +llvm::all_of(deadOps, [&deadOps, &run](Operation *op) { + return allUsesDisappearAfterRewrite(op, deadOps, run); }); } static bool collectLatePureDefinitions(Value value, Operation *runHead, DominanceInfo &dominance, DenseSet &clonedOps) { - if (dominance.dominates(value, runHead)) + if (dominance.dominates(value, runHead)) { return true; + } Operation *defOp = value.getDefiningOp(); - if (!defOp || defOp->getBlock() != runHead->getBlock() || !isPure(defOp)) + if (!defOp || defOp->getBlock() != runHead->getBlock() || !isPure(defOp)) { return false; - if (!clonedOps.insert(defOp).second) + } + if (!clonedOps.insert(defOp).second) { return true; + } return llvm::all_of(defOp->getOperands(), [&](Value operand) { return collectLatePureDefinitions(operand, runHead, dominance, clonedOps); }); @@ -2062,8 +2241,9 @@ static bool isStepMaterializationCostNeutral( SmallVector leaves; collectLeaves(atom, leaves); for (Value leaf : leaves) - if (!collectLatePureDefinitions(leaf, first->op, dominance, clonedOps)) + if (!collectLatePureDefinitions(leaf, first->op, dominance, clonedOps)) { return false; + } } else if (atom->kind == StrideExpr::Kind::Leaf && !collectLatePureDefinitions(atom->leaf, first->op, dominance, clonedOps)) { @@ -2072,14 +2252,15 @@ static bool isStepMaterializationCostNeutral( DenseSet disappearing = deadOffsetOps; disappearing.insert(clonedOps.begin(), clonedOps.end()); - return llvm::all_of(clonedOps, [&](Operation *op) { + return llvm::all_of(clonedOps, [&disappearing, &run](Operation *op) { return allUsesDisappearAfterRewrite(op, disappearing, run); }); } static bool isProfitableDynamicBaseRun(const SequentialRun &run) { - if (!run.stepForm.terms.empty()) + if (!run.stepForm.terms.empty()) { return false; + } unsigned pointerCost = run.candidates.size() - 1; return countDeadDynamicAddPtrs(run) > pointerCost + initialPointerCost(run); @@ -2103,12 +2284,14 @@ static bool isProfitableDirectSymbolicLeafRun( auto firstOffset = firstStrideOperand ? getConstantIntValue(firstStrideOperand) : std::optional(0); - if (!firstOffset || *firstOffset != 0) + if (!firstOffset || *firstOffset != 0) { return false; + } DenseSet deadOffsetOps; - if (!cumulativeOffsetChainDefinitelyDies(run, deadOffsetOps)) + if (!cumulativeOffsetChainDefinitelyDies(run, deadOffsetOps)) { return false; + } return isStepMaterializationCostNeutral(run, deadOffsetOps, dominance); } @@ -2146,21 +2329,24 @@ static void processSequentialBlock(Block *block, DominanceInfo &dominance, for (Operation &op : *block) { originalOps.push_back(&op); const PostUpdateOpInfo *info = getPostUpdateInfo(&op); - if (!info || isAlreadyPostUpdate(&op, *info)) + if (!info || isAlreadyPostUpdate(&op, *info)) { continue; + } Value base, strideOperand; extractBaseAndStrideOperand(&op, *info, base, strideOperand); auto elemBytes = addPtrUnitBytes(base); - if (!elemBytes) + if (!elemBytes) { continue; + } auto unitBytes = strideUnitBytes(&op, info->strideUnit, *elemBytes); - if (!unitBytes) + if (!unitBytes) { continue; + } NormalizedBase normalized = normalizeSequentialBase(base, *elemBytes, exprCache); - auto bucketIt = llvm::find_if(buckets, [&](const SequentialBucket &bucket) { + auto bucketIt = llvm::find_if(buckets, [&op, &normalized](const SequentialBucket &bucket) { return bucket.opName == op.getName().getStringRef() && bucket.rootBase == normalized.root; }); @@ -2190,16 +2376,18 @@ static void processSequentialBlock(Block *block, DominanceInfo &dominance, while (end < candidates.size()) { auto nextStep = analyzeSequentialStep(candidates[end - 1], candidates[end]); - if (!nextStep || !equalAffineForms(firstStep->form, nextStep->form)) + if (!nextStep || !equalAffineForms(firstStep->form, nextStep->form)) { break; + } ++end; } SequentialRun run; run.step = firstStep->expr; run.stepForm = firstStep->form; - for (size_t i = start; i < end; ++i) + for (size_t i = start; i < end; ++i) { run.candidates.push_back(&candidates[i]); + } if (validateSequentialRun(run, dominance) && isProfitableSequentialRun(run, dominance)) { runs.push_back(std::move(run)); @@ -2214,8 +2402,9 @@ static void processSequentialBlock(Block *block, DominanceInfo &dominance, } } - if (runs.empty()) + if (runs.empty()) { return; + } // Materialize every accepted run before erasing any candidate op. for (SequentialRun &run : runs) { @@ -2239,15 +2428,17 @@ static void processSequentialBlock(Block *block, DominanceInfo &dominance, DenseMap opToRun; for (auto [runIdx, run] : llvm::enumerate(runs)) - for (SequentialCandidate *candidate : run.candidates) + for (SequentialCandidate *candidate : run.candidates) { opToRun[candidate->op] = runIdx; + } // Rewrite in original program order so interleaved buckets maintain separate // pointer chains without invalidating one another. for (Operation *op : originalOps) { auto it = opToRun.find(op); - if (it == opToRun.end()) + if (it == opToRun.end()) { continue; + } SequentialRun &run = runs[it->second]; const PostUpdateOpInfo *info = getPostUpdateInfo(op); builder.setInsertionPoint(op); @@ -2256,10 +2447,12 @@ static void processSequentialBlock(Block *block, DominanceInfo &dominance, run.zeroStride, builder) : createPostUpdateOp(op, *info, run.currentPtr, run.strideValue, builder); - for (unsigned result = 0; result < op->getNumResults(); ++result) + for (unsigned result = 0; result < op->getNumResults(); ++result) { op->getResult(result).replaceAllUsesWith(newOp->getResult(result)); - if (!isLast) + } + if (!isLast) { run.currentPtr = newOp->getResult(newOp->getNumResults() - 1); + } op->erase(); } } @@ -2278,7 +2471,7 @@ struct VPTOSoftPostUpdatePass OpBuilder builder(&getContext()); module.walk( - [&](pto::VecScopeOp vecscope) { processVecScope(vecscope, builder); }); + [this, &builder](pto::VecScopeOp vecscope) { processVecScope(vecscope, builder); }); } private: @@ -2287,14 +2480,15 @@ struct VPTOSoftPostUpdatePass // post-order, so nested loops already come before the loops enclosing // them. SmallVector forOps; - vecscope.walk([&](scf::ForOp forOp) { forOps.push_back(forOp); }); + vecscope.walk([&forOps](scf::ForOp forOp) { forOps.push_back(forOp); }); // Process inner-to-outer, i.e. in collection order. The order is load // bearing: rewriting a loop erases it, which also destroys every loop // nested inside it. Visiting an enclosing loop first would leave the // already-collected inner ForOp handles dangling. - for (scf::ForOp forOp : forOps) + for (scf::ForOp forOp : forOps) { processForOp(forOp, builder); + } // Loop rewriting rebuilds ForOps, so collect blocks only after every loop // handle has been consumed. This second phase includes loop bodies and @@ -2302,8 +2496,9 @@ struct VPTOSoftPostUpdatePass SmallVector blocks; collectNestedBlocks(vecscope, vecscope, blocks); DominanceInfo dominance(vecscope->getParentOp()); - for (Block *block : blocks) + for (Block *block : blocks) { processSequentialBlock(block, dominance, builder); + } } void processForOp(scf::ForOp forOp, OpBuilder &builder) { @@ -2314,12 +2509,15 @@ struct VPTOSoftPostUpdatePass for (Operation &op : *forOp.getBody()) { const PostUpdateOpInfo *info = getPostUpdateInfo(&op); - if (!info) + if (!info) { continue; - if (isAlreadyPostUpdate(&op, *info)) + } + if (isAlreadyPostUpdate(&op, *info)) { continue; - if (!isDirectlyInForBody(&op, forOp)) + } + if (!isDirectlyInForBody(&op, forOp)) { continue; + } Value base, strideOperand; extractBaseAndStrideOperand(&op, *info, base, strideOperand); @@ -2330,11 +2528,13 @@ struct VPTOSoftPostUpdatePass // op's lowering expects. Bail on pointers whose addptr unit we cannot // pin down rather than guess at the scale. std::optional elemBytes = addPtrUnitBytes(base); - if (!elemBytes) + if (!elemBytes) { continue; + } auto unitBytes = strideUnitBytes(&op, info->strideUnit, *elemBytes); - if (!unitBytes) + if (!unitBytes) { continue; + } // Analyze each operand independently: accumulator (iter_arg) first, // delta (IV/affine) fallback. Both return a symbolic per-iteration @@ -2345,37 +2545,43 @@ struct VPTOSoftPostUpdatePass strideOperand ? getStride(strideOperand, forOp, deltaCache) : makeConst(0); - if (!deltaBase || !deltaOffset) + if (!deltaBase || !deltaOffset) { continue; + } StrideExprRef total = combineStride(deltaBase, deltaOffset, *elemBytes, *unitBytes); - if (!total) + if (!total) { continue; + } // Reject expressions whose subterms demand conflicting types, or whose // dynamic result cannot be materialized exactly as the op's declared // stride operand type. Type exprResultType; - if (!exprType(total, exprResultType)) + if (!exprType(total, exprResultType)) { continue; + } Type strideType = strideOperand ? strideOperand.getType() : builder.getIndexType(); - if (exprResultType && exprResultType != strideType) + if (exprResultType && exprResultType != strideType) { continue; + } // Reject strides whose constants do not fit the target operand type. - if (!constantsFitType(total, strideType)) + if (!constantsFitType(total, strideType)) { continue; - if (!satisfiesStrideConstraint(total, info->strideConstraint)) + } + if (!satisfiesStrideConstraint(total, info->strideConstraint)) { continue; + } // A stride built only from loop-invariant leaves is materialized before // the loop; otherwise it goes immediately before the candidate op. SmallVector leaves; collectLeaves(total, leaves); bool allInvariant = llvm::all_of( - leaves, [&](Value l) { return forOp.isDefinedOutsideOfLoop(l); }); + leaves, [&forOp](Value l) { return forOp.isDefinedOutsideOfLoop(l); }); StrideExprRef finalExpr = total; if (!allInvariant) { @@ -2399,15 +2605,17 @@ struct VPTOSoftPostUpdatePass Value initPtr = computeInitialPtr( base, initialOffsetOperand, info->strideUnit, *elemBytes, *unitBytes, forOp, builder); - if (!initPtr) + if (!initPtr) { continue; + } rewrites.push_back( {&op, base, strideOperand, strideNew, initPtr, *unitBytes}); } - if (!rewrites.empty()) + if (!rewrites.empty()) { applyPostUpdateRewrites(forOp, rewrites, builder); + } } }; diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 4fb41199da..a46f55f308 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -6674,7 +6674,6 @@ def wait_flag(src: str, dst: str, *, event_id: int = 0): def reserve_buffer(name, *, size, location, auto=True, base=None): - """``pto.reserve_buffer(name, size, location, auto=True, base=None)``.""" space = _normalize_address_space(location) if space not in (_pto.AddressSpace.VEC, _pto.AddressSpace.MAT): raise ValueError( @@ -6691,7 +6690,6 @@ def reserve_buffer(name, *, size, location, auto=True, base=None): def import_reserved_buffer(name, *, peer_func): - """``pto.import_reserved_buffer(name, peer_func=...)``.""" if not isinstance(peer_func, str): spec = getattr(peer_func, "spec", None) role = getattr(spec, "role", None) diff --git a/tools/ptoas/NativeModule.cpp b/tools/ptoas/NativeModule.cpp index 0f0bda26d4..950fe8004b 100644 --- a/tools/ptoas/NativeModule.cpp +++ b/tools/ptoas/NativeModule.cpp @@ -135,8 +135,9 @@ int runPTOASFromPython(const std::vector &arguments) { std::vector storage = arguments; std::vector argv; argv.reserve(storage.size()); - for (std::string &argument : storage) + for (std::string &argument : storage) { argv.push_back(argument.data()); + } py::object contextOwner = py::module_::import("ptoas.mlir.ir").attr("Context")(); diff --git a/tools/ptoas/ObjectEmission.cpp b/tools/ptoas/ObjectEmission.cpp index 806b0cfc9f..3f7cec3216 100644 --- a/tools/ptoas/ObjectEmission.cpp +++ b/tools/ptoas/ObjectEmission.cpp @@ -125,27 +125,32 @@ static std::string sanitizeModuleId(llvm::StringRef raw) { std::string out; out.reserve(raw.size()); for (char c : raw) { - if (std::isalnum(static_cast(c)) || c == '_') + if (std::isalnum(static_cast(c)) || c == '_') { out.push_back(c); - else + } + else { out.push_back('_'); + } } - if (out.empty()) + if (out.empty()) { out = "ptoas_fatobj"; + } return out; } static std::optional getAscendHomePath() { const char *env = std::getenv("ASCEND_HOME_PATH"); - if (!env || !*env) + if (!env || !*env) { return std::nullopt; + } return std::string(env); } static std::optional getEnvPath(llvm::StringRef name) { const char *env = std::getenv(name.str().c_str()); - if (!env || !*env) + if (!env || !*env) { return std::nullopt; + } return std::string(env); } @@ -157,8 +162,9 @@ static std::string joinPath(llvm::StringRef lhs, llvm::StringRef rhs) { static std::optional parseCANNVersionInfo(llvm::StringRef path) { auto buffer = llvm::MemoryBuffer::getFile(path); - if (!buffer) + if (!buffer) { return std::nullopt; + } llvm::StringRef content = buffer.get()->getBuffer(); llvm::SmallVector lines; content.split(lines, '\n'); @@ -166,13 +172,16 @@ static std::optional parseCANNVersionInfo(llvm::StringRef path) { line = line.trim(); llvm::StringRef keys[] = {"Version=", "version="}; for (llvm::StringRef key : keys) { - if (!line.starts_with(key)) + if (!line.starts_with(key)) { continue; + } llvm::StringRef value = line.drop_front(key.size()).trim(); - if (value.consume_front("\"")) + if (value.consume_front("\"")) { value.consume_back("\""); - if (!value.empty()) + } + if (!value.empty()) { return value.str(); + } } } return std::nullopt; @@ -188,18 +197,21 @@ discoverCANNVersion(llvm::StringRef ascendHome) { "aarch64-linux/ascend_all_cann_install.info", "ascend_toolkit_install.info", "ascend_all_cann_install.info", "opp/version.info"}) { - if (auto version = parseCANNVersionInfo(joinPath(ascendHome, relPath))) + if (auto version = parseCANNVersionInfo(joinPath(ascendHome, relPath))) { return version; + } } return std::nullopt; } static std::optional locateProgram(llvm::StringRef envPath, llvm::StringRef fallbackName) { - if (!envPath.empty() && llvm::sys::fs::exists(envPath)) + if (!envPath.empty() && llvm::sys::fs::exists(envPath)) { return envPath.str(); - if (auto found = llvm::sys::findProgramByName(fallbackName)) + } + if (auto found = llvm::sys::findProgramByName(fallbackName)) { return *found; + } return std::nullopt; } @@ -209,24 +221,29 @@ static bool hasPTOISAHeader(llvm::StringRef includeDir) { static void addExistingIncludeDir(llvm::SmallVectorImpl &dirs, llvm::StringRef path) { - if (path.empty() || !llvm::sys::fs::is_directory(path)) + if (path.empty() || !llvm::sys::fs::is_directory(path)) { return; - if (llvm::is_contained(dirs, path)) + } + if (llvm::is_contained(dirs, path)) { return; + } dirs.push_back(path.str()); } static void addPTOISAIncludeDirs(llvm::SmallVectorImpl &dirs, llvm::StringRef ptoIsaPath) { - if (ptoIsaPath.empty() || !llvm::sys::fs::is_directory(ptoIsaPath)) + if (ptoIsaPath.empty() || !llvm::sys::fs::is_directory(ptoIsaPath)) { return; + } std::string includeDir = joinPath(ptoIsaPath, "include"); - if (hasPTOISAHeader(includeDir)) + if (hasPTOISAHeader(includeDir)) { addExistingIncludeDir(dirs, includeDir); + } std::string commonDir = joinPath(ptoIsaPath, "tests/common"); addExistingIncludeDir(dirs, commonDir); - if (hasPTOISAHeader(ptoIsaPath)) + if (hasPTOISAHeader(ptoIsaPath)) { addExistingIncludeDir(dirs, ptoIsaPath); + } } static llvm::SmallVector @@ -234,8 +251,9 @@ discoverCppIncludeDirs(llvm::StringRef ascendHome, llvm::raw_ostream &diagOS, std::string &ptoIsaPath) { llvm::SmallVector includeDirs; - if (auto env = getEnvPath("PTO_ISA_PATH")) + if (auto env = getEnvPath("PTO_ISA_PATH")) { ptoIsaPath = *env; + } else if (auto env = getEnvPath("PTO_ISA_ROOT")) ptoIsaPath = *env; @@ -292,8 +310,9 @@ static std::string resolveTargetCPU(llvm::Module &module, for (llvm::Function &f : module) { if (f.hasFnAttribute("target-cpu")) { std::string cpu = f.getFnAttribute("target-cpu").getValueAsString().str(); - if (!cpu.empty()) + if (!cpu.empty()) { return cpu; + } } } return getTargetCPU(fallback).str(); @@ -305,28 +324,34 @@ class VPTOFatobjArtifacts { : tempFiles(tempFiles) {} bool emitStubSource(StringRef stubSource, llvm::raw_ostream &diagOS) { - if (failed(tempFiles.create("ptoas-host-stub", ".cpp", stubPath, diagOS))) + if (failed(tempFiles.create("ptoas-host-stub", ".cpp", stubPath, diagOS))) { return false; - if (!writeTextFile(stubPath, stubSource, diagOS)) + } + if (!writeTextFile(stubPath, stubSource, diagOS)) { return false; + } return true; } bool initCommandLogs(llvm::raw_ostream &diagOS) { - if (failed(tempFiles.create("ptoas-stderr", ".log", stderrPath, diagOS))) + if (failed(tempFiles.create("ptoas-stderr", ".log", stderrPath, diagOS))) { return false; + } return true; } bool emitCubeObject(llvm::Module *module, const mlir::pto::CANNToolchain &toolchain, llvm::raw_ostream &diagOS) { - if (!module) + if (!module) { return true; - if (failed(tempFiles.create("ptoas-device", ".ll", cubeLLPath, diagOS))) + } + if (failed(tempFiles.create("ptoas-device", ".ll", cubeLLPath, diagOS))) { return false; - if (failed(tempFiles.create("ptoas-device", ".o", cubeObjPath, diagOS))) + } + if (failed(tempFiles.create("ptoas-device", ".o", cubeObjPath, diagOS))) { return false; + } return succeeded(mlir::pto::emitVPTOCubeDeviceObject( *module, cubeLLPath, cubeObjPath, toolchain, stderrPath, diagOS)); } @@ -335,10 +360,12 @@ class VPTOFatobjArtifacts { const mlir::pto::CANNToolchain &toolchain, mlir::pto::VFSIMTSizeFixMode vfsimtSizeFixMode, llvm::raw_ostream &diagOS) { - if (!module) + if (!module) { return true; - if (failed(tempFiles.create("ptoas-device", ".ll", vectorLLPath, diagOS))) + } + if (failed(tempFiles.create("ptoas-device", ".ll", vectorLLPath, diagOS))) { return false; + } std::string rawVectorObjPath; if (failed(tempFiles.create("ptoas-device-vector-raw", ".o", rawVectorObjPath, diagOS))) @@ -360,8 +387,9 @@ class VPTOFatobjArtifacts { mlir::pto::verifyAndPatchVFSIMTSize( *module, rawVectorObjPath, patchedVectorObjPath, vfsimtSizeFixMode, diagOS); - if (failed(result)) + if (failed(result)) { return false; + } vectorObjPath = std::move(result->objectPath); return true; } @@ -369,10 +397,12 @@ class VPTOFatobjArtifacts { bool mergeDeviceObjects(const mlir::pto::CANNToolchain &toolchain, llvm::raw_ostream &diagOS) { llvm::SmallVector deviceObjPaths; - if (!cubeObjPath.empty()) + if (!cubeObjPath.empty()) { deviceObjPaths.push_back(cubeObjPath); - if (!vectorObjPath.empty()) + } + if (!vectorObjPath.empty()) { deviceObjPaths.push_back(vectorObjPath); + } if (deviceObjPaths.empty()) { diagOS << "Error: VPTO fatobj emission requires at least one device module.\n"; return false; @@ -454,8 +484,9 @@ static bool runCommandWithStderr(llvm::StringRef program, std::optional stdinPath) { llvm::SmallVector args; args.reserve(ownedArgs.size()); - for (const std::string &arg : ownedArgs) + for (const std::string &arg : ownedArgs) { args.push_back(arg); + } llvm::SmallVector, 3> redirects = { stdinPath, stderrPath, stderrPath}; @@ -463,18 +494,22 @@ static bool runCommandWithStderr(llvm::StringRef program, bool execFailed = false; int rc = llvm::sys::ExecuteAndWait(program, args, std::nullopt, redirects, 0, 0, &execErr, &execFailed); - if (!execFailed && rc == 0) + if (!execFailed && rc == 0) { return true; + } diagOS << "Error: " << what << " failed\n"; diagOS << "Command:"; - for (llvm::StringRef arg : args) + for (llvm::StringRef arg : args) { diagOS << " " << arg; + } diagOS << "\n"; - if (!execErr.empty()) + if (!execErr.empty()) { diagOS << execErr << "\n"; - if (auto buffer = llvm::MemoryBuffer::getFile(stderrPath)) + } + if (auto buffer = llvm::MemoryBuffer::getFile(stderrPath)) { diagOS << buffer.get()->getBuffer() << "\n"; + } return false; } @@ -558,8 +593,9 @@ static bool compileCppDeviceSourceToObject( "-std=c++17", "-dc", }; - for (const std::string &includeDir : toolchain.cppIncludeDirs) + for (const std::string &includeDir : toolchain.cppIncludeDirs) { args.push_back("-I" + includeDir); + } args.push_back("-c"); args.push_back(cppPath.str()); args.push_back("-o"); @@ -598,8 +634,9 @@ static bool compileCppDeviceSourceToFatobj( "-dc", "-c", }; - for (const std::string &includeDir : toolchain.cppIncludeDirs) + for (const std::string &includeDir : toolchain.cppIncludeDirs) { args.push_back("-I" + includeDir); + } args.push_back(cppPath.str()); args.push_back("-o"); args.push_back(outObjPath.str()); @@ -610,14 +647,17 @@ static bool compileCppDeviceSourceToFatobj( static std::string resolveHostTargetCPU() { if (const char *envCPU = std::getenv("PTOAS_HOST_TARGET_CPU")) { - if (envCPU[0] != '\0') + if (envCPU[0] != '\0') { return std::string(envCPU); + } } std::string hostCPU = llvm::sys::getHostCPUName().str(); - if (hostCPU == "cortex-x925") + if (hostCPU == "cortex-x925") { return "tsv200m"; - if (hostCPU == "znver4" || hostCPU == "znver5") + } + if (hostCPU == "znver4" || hostCPU == "znver5") { return "znver3"; + } return hostCPU; } @@ -730,8 +770,9 @@ static bool mergeDeviceObjects(llvm::ArrayRef deviceObjPaths, llvm::StringRef ldLldPath, llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { - if (deviceObjPaths.empty()) + if (deviceObjPaths.empty()) { return false; + } llvm::SmallVector args = { ldLldPath.str(), @@ -740,8 +781,9 @@ static bool mergeDeviceObjects(llvm::ArrayRef deviceObjPaths, "-Ttext", "0", }; - for (const std::string &path : deviceObjPaths) + for (const std::string &path : deviceObjPaths) { args.push_back(path); + } args.push_back("-o"); args.push_back(outObjPath.str()); args.push_back("-r"); @@ -755,8 +797,9 @@ static bool linkFatobjFiles(llvm::ArrayRef fatobjPaths, const mlir::pto::CANNToolchain &toolchain, llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { - if (fatobjPaths.empty()) + if (fatobjPaths.empty()) { return false; + } llvm::SmallVector args = { toolchain.bishengPath, @@ -766,8 +809,9 @@ static bool linkFatobjFiles(llvm::ArrayRef fatobjPaths, "-o", outObjPath.str(), }; - for (const std::string &path : fatobjPaths) + for (const std::string &path : fatobjPaths) { args.push_back(path); + } return runCommandWithStderr(toolchain.bishengPath, args, stderrPath, diagOS, "fatobj link"); @@ -778,8 +822,9 @@ static bool linkFatobjFiles(llvm::ArrayRef fatobjPaths, mlir::pto::TempFileRegistry::~TempFileRegistry() { cleanup(); } void mlir::pto::TempFileRegistry::cleanup() { - for (const std::string &path : paths) + for (const std::string &path : paths) { llvm::sys::fs::remove(path); + } paths.clear(); } @@ -833,8 +878,9 @@ mlir::pto::CANNToolchain::create(llvm::raw_ostream &diagOS) { toolchain.ascendHomePath, diagOS, toolchain.ptoIsaPath); toolchain.cppIncludeDirs.assign(cppIncludeDirs.begin(), cppIncludeDirs.end()); - if (failed(toolchain.validate(diagOS))) + if (failed(toolchain.validate(diagOS))) { return std::nullopt; + } return toolchain; } @@ -915,8 +961,9 @@ mlir::LogicalResult mlir::pto::emitCppVectorDeviceObject( llvm::StringRef cppSource, llvm::StringRef cppPath, llvm::StringRef outObjPath, const CANNToolchain &toolchain, llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { - if (failed(writeCppSource(cppSource, cppPath, diagOS))) + if (failed(writeCppSource(cppSource, cppPath, diagOS))) { return failure(); + } return compileCppToDeviceObject(cppPath, outObjPath, ObjectEmissionDeviceTarget::Vector, toolchain, stderrPath, diagOS); @@ -926,8 +973,9 @@ mlir::LogicalResult mlir::pto::emitCppCubeDeviceObject( llvm::StringRef cppSource, llvm::StringRef cppPath, llvm::StringRef outObjPath, const CANNToolchain &toolchain, llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { - if (failed(writeCppSource(cppSource, cppPath, diagOS))) + if (failed(writeCppSource(cppSource, cppPath, diagOS))) { return failure(); + } return compileCppToDeviceObject(cppPath, outObjPath, ObjectEmissionDeviceTarget::Cube, toolchain, stderrPath, diagOS); @@ -937,8 +985,9 @@ mlir::LogicalResult mlir::pto::emitCppFatobj( llvm::StringRef cppSource, llvm::StringRef cppPath, llvm::StringRef outObjPath, const CANNToolchain &toolchain, llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { - if (failed(writeCppSource(cppSource, cppPath, diagOS))) + if (failed(writeCppSource(cppSource, cppPath, diagOS))) { return failure(); + } return compileCppDeviceSourceToFatobj(cppPath, outObjPath, toolchain, stderrPath, diagOS) ? success() @@ -971,11 +1020,13 @@ static mlir::LogicalResult renameLLVMFunction(llvm::Module &module, llvm::StringRef sourceName, llvm::StringRef abiName, llvm::raw_ostream &diagOS) { - if (sourceName == abiName) + if (sourceName == abiName) { return mlir::success(); + } llvm::Function *function = module.getFunction(sourceName); - if (!function) + if (!function) { return mlir::success(); + } if (llvm::Function *existing = module.getFunction(abiName); existing && existing != function) { diagOS << "Error: cannot rename LLVM symbol '" << sourceName << "' to '" @@ -990,14 +1041,16 @@ static mlir::LogicalResult applyVPTOLLVMABINames(llvm::Module &module, llvm::StringRef suffix, llvm::raw_ostream &diagOS) { for (llvm::Function &function : module) { - if (function.isDeclaration() || !function.hasExternalLinkage()) + if (function.isDeclaration() || !function.hasExternalLinkage()) { continue; + } llvm::StringRef name = function.getName(); if (name.empty() || isVPTOKernelABISymbol(name) || isLegacyVPTOPublicABISymbol(name)) continue; - if (failed(renameLLVMFunction(module, name, (name + suffix).str(), diagOS))) + if (failed(renameLLVMFunction(module, name, (name + suffix).str(), diagOS))) { return mlir::failure(); + } } return mlir::success(); } @@ -1011,8 +1064,9 @@ mlir::LogicalResult mlir::pto::emitVPTOVectorDeviceObject( toolchain.vptoPublicABISuffix(ObjectEmissionDeviceTarget::Vector), diagOS))) return failure(); - if (failed(writeLLVMModule(module, llPath, diagOS))) + if (failed(writeLLVMModule(module, llPath, diagOS))) { return failure(); + } return compileDeviceLLVMToObject(llPath, outObjPath, resolveTargetCPU(module, ObjectEmissionDeviceTarget::Vector), @@ -1030,8 +1084,9 @@ mlir::LogicalResult mlir::pto::emitVPTOCubeDeviceObject( toolchain.vptoPublicABISuffix(ObjectEmissionDeviceTarget::Cube), diagOS))) return failure(); - if (failed(writeLLVMModule(module, llPath, diagOS))) + if (failed(writeLLVMModule(module, llPath, diagOS))) { return failure(); + } return compileDeviceLLVMToObject(llPath, outObjPath, resolveTargetCPU(module, ObjectEmissionDeviceTarget::Cube), @@ -1052,17 +1107,21 @@ mlir::LogicalResult mlir::pto::emitFatobjLLVM( } VPTOFatobjArtifacts artifacts(tempFiles); - if (!artifacts.emitStubSource(stubSource, diagOS)) + if (!artifacts.emitStubSource(stubSource, diagOS)) { return failure(); - if (!artifacts.initCommandLogs(diagOS)) + } + if (!artifacts.initCommandLogs(diagOS)) { return failure(); - if (!artifacts.emitCubeObject(cubeModule, toolchain, diagOS)) + } + if (!artifacts.emitCubeObject(cubeModule, toolchain, diagOS)) { return failure(); + } if (!artifacts.emitVectorObject(vectorModule, toolchain, vfsimtSizeFixMode, diagOS)) return failure(); - if (!artifacts.mergeDeviceObjects(toolchain, diagOS)) + if (!artifacts.mergeDeviceObjects(toolchain, diagOS)) { return failure(); + } constexpr llvm::StringLiteral targetCPU = "dav-c310"; if (!artifacts.compileHostStubToFatobj(toolchain, moduleId, targetCPU, @@ -1114,29 +1173,35 @@ mlir::LogicalResult mlir::pto::emitFatobjLLVMWithRuntime( } std::optional toolchain = CANNToolchain::create(diagOS); - if (!toolchain) + if (!toolchain) { return failure(); + } TempFileRegistry tempFiles; VPTOFatobjArtifacts artifacts(tempFiles); - if (!artifacts.emitStubSource(stubSource, diagOS)) + if (!artifacts.emitStubSource(stubSource, diagOS)) { return failure(); - if (!artifacts.initCommandLogs(diagOS)) + } + if (!artifacts.initCommandLogs(diagOS)) { return failure(); + } - if (!artifacts.emitCubeObject(cubeModule, *toolchain, diagOS)) + if (!artifacts.emitCubeObject(cubeModule, *toolchain, diagOS)) { return failure(); + } if (!artifacts.emitVectorObject(vectorModule, *toolchain, vfsimtSizeFixMode, diagOS)) return failure(); - if (!artifacts.mergeDeviceObjects(*toolchain, diagOS)) + if (!artifacts.mergeDeviceObjects(*toolchain, diagOS)) { return failure(); + } std::string moduleId = sanitizeModuleId(outputFile.getFilename()); constexpr llvm::StringLiteral hostTargetCPU = "dav-c310"; - if (!artifacts.compileHostStub(*toolchain, moduleId, hostTargetCPU, diagOS)) + if (!artifacts.compileHostStub(*toolchain, moduleId, hostTargetCPU, diagOS)) { return failure(); + } if (!artifacts.repackFatObj(*toolchain, moduleId, hostTargetCPU, outputFile.getFilename(), diagOS)) diff --git a/tools/ptoas/VFSIMTSizePatcher.cpp b/tools/ptoas/VFSIMTSizePatcher.cpp index 0785f8ebff..b3450adc01 100644 --- a/tools/ptoas/VFSIMTSizePatcher.cpp +++ b/tools/ptoas/VFSIMTSizePatcher.cpp @@ -112,8 +112,9 @@ template static std::optional takeExpected(llvm::Expected value, llvm::raw_ostream &diagOS, const llvm::Twine &context) { - if (value) + if (value) { return std::move(*value); + } diagOS << "Error: VF_SIMT size patch: " << context << ": " << llvm::toString(value.takeError()) << "\n"; return std::nullopt; @@ -123,8 +124,9 @@ static bool functionContainsInlineAsm(const llvm::Function &function) { for (const llvm::BasicBlock &block : function) { for (const llvm::Instruction &instruction : block) { const auto *call = llvm::dyn_cast(&instruction); - if (call && call->isInlineAsm()) + if (call && call->isInlineAsm()) { return true; + } } } return false; @@ -147,8 +149,9 @@ collectManifest(llvm::Module &module, llvm::raw_ostream &diagOS) { for (llvm::BasicBlock &block : caller) { for (llvm::Instruction &instruction : block) { auto *call = llvm::dyn_cast(&instruction); - if (!call || call->getCallingConv() != llvm::CallingConv::SimtEntry) + if (!call || call->getCallingConv() != llvm::CallingConv::SimtEntry) { continue; + } llvm::Function *callee = call->getCalledFunction(); if (!callee || callee->isDeclaration()) { emitError(diagOS, @@ -188,10 +191,12 @@ readFunctions(const llvm::object::ELFObjectFileBase &object, unsigned symbolTables = 0; for (const llvm::object::SectionRef §ion : object.sections()) { llvm::object::ELFSectionRef elfSection(section); - if (elfSection.getType() == llvm::ELF::SHT_SYMTAB) + if (elfSection.getType() == llvm::ELF::SHT_SYMTAB) { ++symbolTables; - if ((elfSection.getFlags() & llvm::ELF::SHF_EXECINSTR) == 0) + } + if ((elfSection.getFlags() & llvm::ELF::SHF_EXECINSTR) == 0) { continue; + } ++executableSections; } if (symbolTables != 1) { @@ -207,22 +212,26 @@ readFunctions(const llvm::object::ELFObjectFileBase &object, } for (const llvm::object::ELFSymbolRef symbol : object.symbols()) { - if (symbol.getELFType() != llvm::ELF::STT_FUNC) + if (symbol.getELFType() != llvm::ELF::STT_FUNC) { continue; + } auto name = takeExpected(symbol.getName(), diagOS, "failed to read symbol name"); - if (!name) + if (!name) { return failure(); - if (!requiredFunctions.contains(*name)) + } + if (!requiredFunctions.contains(*name)) { continue; + } auto address = takeExpected(symbol.getAddress(), diagOS, llvm::Twine("failed to read address for '") + *name + "'"); auto sectionIt = takeExpected(symbol.getSection(), diagOS, llvm::Twine("failed to read section for '") + *name + "'"); - if (!address || !sectionIt) + if (!address || !sectionIt) { return failure(); + } if (*sectionIt == object.section_end()) { emitError(diagOS, llvm::Twine("function '") + *name + "' is undefined"); return failure(); @@ -263,8 +272,9 @@ static bool isVFSIMT(uint64_t instruction) { static std::optional decodeMOVKChunk(uint32_t instruction, unsigned chunk, unsigned targetRegister) { constexpr uint32_t kNop = 0x41400000; - if (instruction == kNop) + if (instruction == kNop) { return 0; + } constexpr uint32_t kMOVKFixedMask = 0xffc10000; const uint32_t expected = chunk == 1 ? 0x07410000 : 0x07810000; @@ -292,8 +302,9 @@ static std::optional decodeTargetAddress(StringRef bytes, // targetRegister comes from the VF_SIMT encoding. The sequence computes: // target PC = address of MOV PC + sign_extend(relativeWords) * 4. // Reject any other sequence instead of guessing its target. - if (instructionOffset < 24) + if (instructionOffset < 24) { return std::nullopt; + } // Follow the register named by VF_SIMT back through MOVI/MOVK/SHLI/ADD. const unsigned targetRegister = (instruction >> kVFSIMTRegisterShift) & kScalarRegisterMask; @@ -327,13 +338,15 @@ static std::optional decodeTargetAddress(StringRef bytes, const unsigned pcRegister = (movPc >> kScalarDestinationShift) & kScalarRegisterMask; const unsigned addSourceRegister = (add >> 7) & kScalarRegisterMask; - if (pcRegister != addSourceRegister) + if (pcRegister != addSourceRegister) { return std::nullopt; + } std::optional upper1 = decodeMOVKChunk(movk1, 1, targetRegister); std::optional upper2 = decodeMOVKChunk(movk2, 2, targetRegister); - if (!upper1 || !upper2) + if (!upper1 || !upper2) { return std::nullopt; + } // MOVI and the optional MOVK instructions form a signed 48-bit word offset. const uint64_t encodedRelativeWords = (movi & 0xffff) | @@ -346,8 +359,9 @@ static std::optional decodeTargetAddress(StringRef bytes, const int64_t byteOffset = relativeWords * kInstructionBytes; if (byteOffset < 0) { const uint64_t magnitude = static_cast(-byteOffset); - if (magnitude > pcAddress) + if (magnitude > pcAddress) { return std::nullopt; + } return pcAddress - magnitude; } if (static_cast(byteOffset) > @@ -375,8 +389,9 @@ decodeCallSites(const ELFFunction &caller, StringRef objectBytes, offset += kInstructionBytes) { uint64_t instruction = llvm::support::endian::read64le( reinterpret_cast(bytes.data() + offset)); - if (!isVFSIMT(instruction)) + if (!isVFSIMT(instruction)) { continue; + } std::optional target = decodeTargetAddress(bytes, caller.address, offset, instruction); if (!target) { @@ -423,8 +438,9 @@ buildPatchPlan(llvm::ArrayRef manifest, // address. Call order alone is not sufficient proof that a patch is safe. llvm::DenseMap> callsByCaller; - for (const SimtCallSite &call : manifest) + for (const SimtCallSite &call : manifest) { callsByCaller[call.callerName].push_back(&call); + } llvm::SmallVector plan; for (auto &entry : callsByCaller) { @@ -435,8 +451,9 @@ buildPatchPlan(llvm::ArrayRef manifest, return failure(); } auto decoded = decodeCallSites(callerIt->second, objectBytes, diagOS); - if (failed(decoded)) + if (failed(decoded)) { return failure(); + } auto &calls = entry.second; struct ResolvedCall { const SimtCallSite *manifest = nullptr; @@ -478,8 +495,9 @@ buildPatchPlan(llvm::ArrayRef manifest, for (const DecodedCallSite &decodedCall : *decoded) { ResolvedCall *matched = nullptr; for (ResolvedCall &candidate : resolvedCalls) { - if (candidate.callee->address != decodedCall.targetAddress) + if (candidate.callee->address != decodedCall.targetAddress) { continue; + } if (matched && matched->callee->name != candidate.callee->name) { emitError(diagOS, llvm::Twine("caller '") + entry.first + @@ -487,8 +505,9 @@ buildPatchPlan(llvm::ArrayRef manifest, llvm::Twine::utohexstr(decodedCall.targetAddress)); return failure(); } - if (!matched) + if (!matched) { matched = &candidate; + } } if (!matched) { emitError(diagOS, @@ -504,8 +523,9 @@ buildPatchPlan(llvm::ArrayRef manifest, static_cast(matched->callee->size / kInstructionBytes)}); } for (const ResolvedCall &call : resolvedCalls) { - if (call.observed) + if (call.observed) { continue; + } emitError(diagOS, llvm::Twine("caller '") + entry.first + "' has no decoded VF_SIMT callsite for callee '" + call.callee->name + "'"); @@ -532,13 +552,15 @@ validateNoRelocationOverlap(const llvm::object::ELFObjectFileBase &object, // Reject a VF_SIMT covered by a relocation: the linker could overwrite the // patched instruction and invalidate the checks performed on the raw object. for (const llvm::object::SectionRef §ion : object.sections()) { - if (section.relocation_begin() == section.relocation_end()) + if (section.relocation_begin() == section.relocation_end()) { continue; + } auto relocatedSection = takeExpected(section.getRelocatedSection(), diagOS, "failed to identify the section targeted by relocations"); - if (!relocatedSection) + if (!relocatedSection) { return failure(); + } if (*relocatedSection == object.section_end()) { emitError(diagOS, "relocation section has no target section"); return failure(); @@ -592,11 +614,13 @@ analyzeObject(llvm::ArrayRef manifest, StringRef objectPath, emitError(diagOS, "input is not an ELF object"); return failure(); } - if (failed(validateObjectHeader(*object, diagOS))) + if (failed(validateObjectHeader(*object, diagOS))) { return failure(); + } auto functions = readFunctions(*object, manifest, diagOS); - if (failed(functions)) + if (failed(functions)) { return failure(); + } auto plan = buildPatchPlan(manifest, *functions, buffer->getBuffer(), diagOS); if (failed(plan) || failed(validateNoRelocationOverlap(*object, *plan, diagOS))) @@ -627,8 +651,9 @@ static LogicalResult validatePatchedBytes(StringRef rawBytes, } } for (size_t offset = 0; offset < rawBytes.size(); ++offset) { - if (rawBytes[offset] == patchedBytes[offset]) + if (rawBytes[offset] == patchedBytes[offset]) { continue; + } if (!llvm::any_of(plan, [offset](const PatchRecord &record) { return offset >= record.decoded.fileOffset && offset < record.decoded.fileOffset + sizeof(uint64_t); @@ -674,20 +699,23 @@ FailureOr mlir::pto::verifyAndPatchVFSIMTSize( // output when any callsite is unsafe or inconsistent. VFSIMTSizePatchResult result; result.objectPath = rawObjectPath.str(); - if (mode == VFSIMTSizeFixMode::Off) + if (mode == VFSIMTSizeFixMode::Off) { return result; + } auto manifest = collectManifest(module, diagOS); - if (failed(manifest)) + if (failed(manifest)) { return failure(); + } if (manifest->empty()) { diagOS << "PTOAS: VF_SIMT size verification passed; no patch required\n"; return result; } auto analysis = analyzeObject(*manifest, rawObjectPath, diagOS); - if (failed(analysis)) + if (failed(analysis)) { return failure(); + } std::string patchedBytes = analysis->buffer->getBuffer().str(); for (const PatchRecord &record : analysis->plan) { @@ -749,8 +777,9 @@ FailureOr mlir::pto::verifyAndPatchVFSIMTSize( if (failed(validatePatchedBytes(analysis->buffer->getBuffer(), patchedBytes, analysis->plan, diagOS))) return failure(); - if (failed(writePatchedObject(patchedObjectPath, patchedBytes, diagOS))) + if (failed(writePatchedObject(patchedObjectPath, patchedBytes, diagOS))) { return failure(); + } auto writtenBuffer = llvm::MemoryBuffer::getFile(patchedObjectPath); if (!writtenBuffer || writtenBuffer.get()->getBuffer() != patchedBytes) { @@ -787,8 +816,9 @@ FailureOr mlir::pto::verifyAndPatchVFSIMTSize( } for (const PatchRecord &record : analysis->plan) { - if (record.decoded.codeSize != kInvalidVFSIMTSize) + if (record.decoded.codeSize != kInvalidVFSIMTSize) { continue; + } diagOS << "PTOAS: patched VF_SIMT size\n" << " caller: " << record.manifest->callerName << "\n" << " callee: " << record.manifest->calleeName << "\n" diff --git a/tools/ptoas/VPTOHostStubEmission.cpp b/tools/ptoas/VPTOHostStubEmission.cpp index 103e4f80eb..ca4956076f 100644 --- a/tools/ptoas/VPTOHostStubEmission.cpp +++ b/tools/ptoas/VPTOHostStubEmission.cpp @@ -26,16 +26,19 @@ struct VPTOKernelStubDecl { }; static std::string getLogicalKernelName(llvm::StringRef symbol) { - if (symbol.ends_with("_mix_aiv")) + if (symbol.ends_with("_mix_aiv")) { return symbol.drop_back(strlen("_mix_aiv")).str(); - if (symbol.ends_with("_mix_aic")) + } + if (symbol.ends_with("_mix_aic")) { return symbol.drop_back(strlen("_mix_aic")).str(); + } return symbol.str(); } static std::string getStubScalarCType(Type type) { - if (isa(type)) + if (isa(type)) { return "long long"; + } if (auto intType = dyn_cast(type)) { switch (intType.getWidth()) { case 1: @@ -52,18 +55,21 @@ static std::string getStubScalarCType(Type type) { } } if (auto floatType = dyn_cast(type)) { - if (floatType.isF32()) + if (floatType.isF32()) { return "float"; - if (floatType.isF64()) + } + if (floatType.isF64()) { return "double"; + } return "short"; } return "long long"; } static std::string getStubCType(Type type) { - if (isa(type)) + if (isa(type)) { return "__gm__ void *"; + } return getStubScalarCType(type); } @@ -77,14 +83,16 @@ static LogicalResult collectVPTOKernelStubDecls( for (ModuleOp module : modules) { module.walk([&](func::FuncOp func) { - if (!pto::isPTOEntryFunction(func)) + if (!pto::isPTOEntryFunction(func)) { return; + } std::string logicalName = getLogicalKernelName(func.getSymName()); SmallVector argTypes; argTypes.reserve(func.getNumArguments()); - for (Type type : func.getArgumentTypes()) + for (Type type : func.getArgumentTypes()) { argTypes.push_back(getStubCType(type)); + } auto [it, inserted] = logicalNameToIndex.try_emplace(logicalName, decls.size()); @@ -115,8 +123,9 @@ LogicalResult mlir::pto::emitVPTOHostStubSource(ArrayRef modules, std::string &stubSource, llvm::raw_ostream &diagOS) { SmallVector stubDecls; - if (failed(collectVPTOKernelStubDecls(modules, stubDecls, diagOS))) + if (failed(collectVPTOKernelStubDecls(modules, stubDecls, diagOS))) { return failure(); + } if (stubDecls.empty()) { diagOS << "Error: no PTO entry functions found for host stub emission.\n"; @@ -129,8 +138,9 @@ LogicalResult mlir::pto::emitVPTOHostStubSource(ArrayRef modules, for (const VPTOKernelStubDecl &decl : stubDecls) { os << "extern \"C\" __global__ AICORE void " << decl.logicalName << "("; for (size_t i = 0; i < decl.argTypes.size(); ++i) { - if (i) + if (i) { os << ", "; + } os << decl.argTypes[i] << " arg" << i; } os << ") {}\n"; diff --git a/tools/ptoas/driver.cpp b/tools/ptoas/driver.cpp index 06644f4307..2dbae2f76c 100644 --- a/tools/ptoas/driver.cpp +++ b/tools/ptoas/driver.cpp @@ -68,8 +68,9 @@ static bool parseRequestedOutputCANNVersion( llvm::StringRef versionText, std::optional &version, llvm::raw_ostream &diagOS) { version.reset(); - if (versionText.empty()) + if (versionText.empty()) { return true; + } std::optional parsed = mlir::pto::parseCANNVersion(versionText); if (!parsed) { @@ -91,16 +92,18 @@ static bool hasCLIOption(int argc, char **argv, llvm::StringRef option) { const std::string optionWithValue = (option + "=").str(); for (int i = 1; i < argc; ++i) { llvm::StringRef arg(argv[i]); - if (arg == option || arg.starts_with(optionWithValue)) + if (arg == option || arg.starts_with(optionWithValue)) { return true; + } } return false; } static std::string normalizePTOASArch(llvm::StringRef archValue) { std::string normalized = archValue.str(); - for (char &c : normalized) + for (char &c : normalized) { c = static_cast(std::tolower(static_cast(c))); + } return normalized; } @@ -113,8 +116,9 @@ detectPTOASTextualModuleArch(llvm::StringRef text) { llvm::SmallVector matches; llvm::Regex archRegex( R"ptoarch("?(pto\.target_arch)"?[[:space:]]*=[[:space:]]*"([[:alpha:][:digit:]_]+)")ptoarch"); - if (!archRegex.match(text, &matches) || matches.size() < 3) + if (!archRegex.match(text, &matches) || matches.size() < 3) { return std::nullopt; + } return normalizePTOASArch(matches[2]); } @@ -144,10 +148,12 @@ static bool resolveTextInputArch(llvm::StringRef buffer, bool cliArchSpecified, return true; } - if (auto detectedArch = detectPTOASTextualModuleArch(buffer)) + if (auto detectedArch = detectPTOASTextualModuleArch(buffer)) { arch = *detectedArch; - if (!isSupportedPTOASArch(arch)) + } + if (!isSupportedPTOASArch(arch)) { arch = "a3"; + } return true; } @@ -164,8 +170,9 @@ static OwningOpRef decodePTOBCModule(llvm::StringRef buffer, } #else OwningOpRef module = ptobc::decodePTOBCToModule(bytes, context); - if (!module) + if (!module) { llvm::errs() << "Error: Failed to decode PTOBC.\n"; + } return module; #endif } @@ -238,14 +245,16 @@ loadInputModule(std::unique_ptr inputBuffer, if (isSupportedPTOASArch(moduleArch)) { arch = std::move(moduleArch); } else { - if (!isSupportedPTOASArch(arch)) + if (!isSupportedPTOASArch(arch)) { arch = "a3"; + } moduleOp->setAttr("pto.target_arch", mlir::StringAttr::get(moduleOp->getContext(), arch)); } } else { - if (!isSupportedPTOASArch(arch)) + if (!isSupportedPTOASArch(arch)) { arch = "a3"; + } moduleOp->setAttr("pto.target_arch", mlir::StringAttr::get(moduleOp->getContext(), arch)); } @@ -260,8 +269,9 @@ loadInputModule(std::unique_ptr inputBuffer, static bool parseDriverBackend(llvm::StringRef backendStr, mlir::pto::PTOBackend &out) { std::string s = backendStr.str(); - for (char &c : s) + for (char &c : s) { c = static_cast(std::tolower(static_cast(c))); + } if (s == "emitc") { out = mlir::pto::PTOBackend::EmitC; return true; @@ -278,8 +288,9 @@ parseDriverBackendAttr(Operation *op, std::optional &backend) { backend = std::nullopt; Attribute rawBackendAttr = op->getAttr("pto.backend"); - if (!rawBackendAttr) + if (!rawBackendAttr) { return success(); + } auto backendAttr = dyn_cast(rawBackendAttr); if (!backendAttr) { @@ -299,8 +310,9 @@ parseDriverBackendAttr(Operation *op, static bool isBackendPartitionedContainer(ModuleOp module) { Block *body = module.getBody(); - if (!body) + if (!body) { return false; + } return llvm::all_of(body->getOperations(), [](Operation &op) { return isa(op); }); } @@ -333,11 +345,13 @@ static SmallVector collectDirectCalleeNames(ModuleOp module) { static SmallVector collectDirectCalleeNames(func::FuncOp funcOp) { SmallVector names; - if (!funcOp || funcOp.isDeclaration()) + if (!funcOp || funcOp.isDeclaration()) { return names; + } funcOp.walk([&](func::CallOp callOp) { - if (callOp->getParentOfType() != funcOp) + if (callOp->getParentOfType() != funcOp) { return; + } names.push_back(callOp.getCalleeAttr().getLeafReference()); }); llvm::sort(names); @@ -358,8 +372,9 @@ static void copyModuleAttrsToJobModule(ModuleOp source, ModuleOp jobModule) { static func::FuncOp findFunctionByLogicalName(ModuleOp module, StringRef logicalName) { for (func::FuncOp funcOp : module.getOps()) { - if (funcOp.getSymName() == logicalName) + if (funcOp.getSymName() == logicalName) { return funcOp; + } } return {}; } @@ -372,19 +387,23 @@ static func::FuncOp findFunctionBySymbolName(ModuleOp module, static func::FuncOp findFunctionForPeerReference(ModuleOp module, StringRef peerRef) { - if (func::FuncOp exact = findFunctionBySymbolName(module, peerRef)) + if (func::FuncOp exact = findFunctionBySymbolName(module, peerRef)) { return exact; + } func::FuncOp privateMatch; for (func::FuncOp funcOp : module.getOps()) { - if (mlir::pto::getPTODSLLogicalNameOrSymbolName(funcOp) != peerRef) + if (mlir::pto::getPTODSLLogicalNameOrSymbolName(funcOp) != peerRef) { continue; + } auto visibility = funcOp->getAttrOfType("sym_visibility"); - if (!visibility || visibility.getValue() != "private") + if (!visibility || visibility.getValue() != "private") { return funcOp; - if (!privateMatch) + } + if (!privateMatch) { privateMatch = funcOp; + } } return privateMatch; } @@ -396,12 +415,14 @@ findSiblingSourceFunction(ModuleOp outer, ModuleOp targetChild, SmallVector exactMatches; SmallVector logicalMatches; for (ModuleOp child : outer.getOps()) { - if (child == targetChild) + if (child == targetChild) { continue; + } for (func::FuncOp funcOp : child.getOps()) { auto visibility = funcOp->getAttrOfType("sym_visibility"); - if (visibility && visibility.getValue() == "private") + if (visibility && visibility.getValue() == "private") { continue; + } if (funcOp.getSymName() == symbolName) { exactMatches.push_back(funcOp); continue; @@ -418,8 +439,9 @@ findSiblingSourceFunction(ModuleOp outer, ModuleOp targetChild, << "'; found multiple sibling public func.func definitions"; return failure(); } - if (!exactMatches.empty()) + if (!exactMatches.empty()) { return exactMatches.front(); + } if (logicalMatches.size() > 1) { targetChild.emitError("mixed-backend child assembly does not yet support ambiguous cross-child logical ") @@ -427,16 +449,18 @@ findSiblingSourceFunction(ModuleOp outer, ModuleOp targetChild, << "'; found multiple sibling public func.func definitions"; return failure(); } - if (!logicalMatches.empty()) + if (!logicalMatches.empty()) { return logicalMatches.front(); + } return func::FuncOp(); } static LogicalResult verifyImportedPeerCloneContract(func::FuncOp peerSource, StringRef logicalName) { SmallVector directCalleeNames = collectDirectCalleeNames(peerSource); - if (directCalleeNames.empty()) + if (directCalleeNames.empty()) { return success(); + } peerSource.emitError( "mixed-backend child assembly does not yet support transitive cross-child function closure for imported peer '") @@ -468,18 +492,21 @@ static func::FuncOp cloneFunctionDeclarationIntoModule(ModuleOp jobModule, StringRef visibility) { func::FuncOp cloned = cloneFunctionIntoModule(jobModule, sourceFunc, newName, visibility); - while (!cloned.getBody().empty()) + while (!cloned.getBody().empty()) { cloned.getBody().front().erase(); + } return cloned; } static void rewriteExportedFunctionToLogicalWrapper(func::FuncOp exportedFunc, StringRef logicalName) { - if (logicalName == exportedFunc.getSymName()) + if (logicalName == exportedFunc.getSymName()) { return; + } - while (!exportedFunc.getBody().empty()) + while (!exportedFunc.getBody().empty()) { exportedFunc.getBody().front().erase(); + } Block *entry = exportedFunc.addEntryBlock(); OpBuilder builder(entry, entry->begin()); @@ -497,16 +524,18 @@ verifyInChildLogicalWrapperAmbiguity(ModuleOp targetChild, auto kernelKindAttr = exportedFunc->getAttrOfType( mlir::pto::FunctionKernelKindAttr::name); - if (kernelKindAttr) + if (kernelKindAttr) { continue; + } StringRef logicalName = mlir::pto::getPTODSLLogicalNameOrSymbolName(exportedFunc); grouped[logicalName].push_back(exportedFunc); } for (const auto &entry : grouped) { - if (entry.second.size() <= 1) + if (entry.second.size() <= 1) { continue; + } targetChild.emitError( "mixed-backend child assembly does not yet support ambiguous in-child logical reference '@") << entry.first @@ -528,14 +557,16 @@ buildBackendChildCompileUnit(ModuleOp outer, ModuleOp targetChild) { SmallVector directCalleeNames = collectDirectCalleeNames(targetChild); for (StringRef calleeName : directCalleeNames) { - if (findFunctionByLogicalName(jobModule, calleeName)) + if (findFunctionByLogicalName(jobModule, calleeName)) { continue; + } FailureOr siblingSourceOr = findSiblingSourceFunction(outer, targetChild, calleeName, /*allowLogicalNameMatch=*/false, /*referenceKind=*/"function reference"); - if (failed(siblingSourceOr)) + if (failed(siblingSourceOr)) { return failure(); + } func::FuncOp siblingSource = *siblingSourceOr; if (!siblingSource) { targetChild.emitError( @@ -552,22 +583,26 @@ buildBackendChildCompileUnit(ModuleOp outer, ModuleOp targetChild) { SmallVector exportedFuncs; for (func::FuncOp funcOp : jobModule.getOps()) { auto visibility = funcOp->getAttrOfType("sym_visibility"); - if (visibility && visibility.getValue() == "private") + if (visibility && visibility.getValue() == "private") { continue; - if (funcOp.isExternal()) + } + if (funcOp.isExternal()) { continue; + } exportedFuncs.push_back(funcOp); } - if (failed(verifyInChildLogicalWrapperAmbiguity(targetChild, exportedFuncs))) + if (failed(verifyInChildLogicalWrapperAmbiguity(targetChild, exportedFuncs))) { return failure(); + } for (func::FuncOp exportedFunc : exportedFuncs) { StringRef logicalName = mlir::pto::getPTODSLLogicalNameOrSymbolName(exportedFunc); auto kernelKindAttr = exportedFunc->getAttrOfType( mlir::pto::FunctionKernelKindAttr::name); - if (kernelKindAttr) + if (kernelKindAttr) { continue; + } if (!findFunctionByLogicalName(jobModule, logicalName)) { cloneFunctionIntoModule(jobModule, exportedFunc, logicalName, "private"); } @@ -579,8 +614,9 @@ buildBackendChildCompileUnit(ModuleOp outer, ModuleOp targetChild) { findSiblingSourceFunction(outer, targetChild, logicalName, /*allowLogicalNameMatch=*/true, /*referenceKind=*/"peer_func reference"); - if (failed(peerSourceOr)) + if (failed(peerSourceOr)) { return failure(); + } func::FuncOp peerSource = *peerSourceOr; if (!peerSource) { targetChild.emitError( @@ -589,16 +625,19 @@ buildBackendChildCompileUnit(ModuleOp outer, ModuleOp targetChild) { << "'; each import_reserved_buffer peer_func must resolve to one sibling public func.func"; return failure(); } - if (failed(verifyImportedPeerCloneContract(peerSource, logicalName))) + if (failed(verifyImportedPeerCloneContract(peerSource, logicalName))) { return failure(); + } StringRef peerSymbolName = peerSource.getSymName(); - if (!findFunctionBySymbolName(jobModule, peerSymbolName)) + if (!findFunctionBySymbolName(jobModule, peerSymbolName)) { cloneFunctionIntoModule(jobModule, peerSource, peerSymbolName, "private"); + } jobModule.walk([&](pto::ImportReservedBufferOp importOp) { - if (importOp.getPeerFuncAttr().getValue() != logicalName) + if (importOp.getPeerFuncAttr().getValue() != logicalName) { return; + } importOp.setPeerFuncAttr( FlatSymbolRefAttr::get(jobModule.getContext(), peerSymbolName)); }); @@ -607,10 +646,12 @@ buildBackendChildCompileUnit(ModuleOp outer, ModuleOp targetChild) { jobModule.walk([&](pto::ImportReservedBufferOp importOp) { StringRef peerRef = importOp.getPeerFuncAttr().getValue(); func::FuncOp localPeer = findFunctionForPeerReference(jobModule, peerRef); - if (!localPeer) + if (!localPeer) { return; - if (localPeer.getSymName() == peerRef) + } + if (localPeer.getSymName() == peerRef) { return; + } importOp.setPeerFuncAttr( FlatSymbolRefAttr::get(jobModule.getContext(), localPeer.getSymName())); }); @@ -626,26 +667,31 @@ static std::string summarizeMixedChildModule(ModuleOp module) { std::string summary; llvm::raw_string_ostream os(summary); - if (auto backendAttr = module->getAttrOfType("pto.backend")) + if (auto backendAttr = module->getAttrOfType("pto.backend")) { os << "backend=" << backendAttr.getValue() << " "; - if (auto kindAttr = module->getAttrOfType("pto.kernel_kind")) + } + if (auto kindAttr = module->getAttrOfType("pto.kernel_kind")) { os << "kernel_kind=" << kindAttr.getValue() << " "; + } SmallVector exportedNames; for (func::FuncOp funcOp : module.getOps()) { auto visibility = funcOp->getAttrOfType("sym_visibility"); - if (visibility && visibility.getValue() == "private") + if (visibility && visibility.getValue() == "private") { continue; - if (funcOp.isExternal()) + } + if (funcOp.isExternal()) { continue; + } exportedNames.push_back(funcOp.getSymName().str()); } if (!exportedNames.empty()) { os << "exports=["; for (size_t i = 0; i < exportedNames.size(); ++i) { - if (i) + if (i) { os << ", "; + } os << exportedNames[i]; } os << "]"; @@ -693,8 +739,9 @@ mlir::pto::PTOASContext::~PTOASContext() = default; LogicalResult mlir::pto::PTOASContext::initializeEnvironment(bool requiresToolchain, llvm::raw_ostream &diagOS) { - if (requiresToolchain) + if (requiresToolchain) { return initializeToolchain(diagOS); + } return success(); } @@ -750,12 +797,14 @@ std::string mlir::pto::PTOASContext::allocModuleId() { LogicalResult mlir::pto::PTOASContext::initializeToolchain(llvm::raw_ostream &diagOS) { - if (toolchain) + if (toolchain) { return success(); + } std::optional discovered = mlir::pto::CANNToolchain::create(diagOS); - if (!discovered) + if (!discovered) { return failure(); + } std::optional parsedVersion = parseCANNVersion(discovered->cannVersionString); if (!parsedVersion) { @@ -870,12 +919,14 @@ class EmitCBackendChildJob final : public BackendChildJob { } std::string fatobjPath; - if (failed(context.createTempPath("ptoas-emitc-fatobj", ".o", fatobjPath))) + if (failed(context.createTempPath("ptoas-emitc-fatobj", ".o", fatobjPath))) { return failure(); + } const mlir::pto::CANNToolchain *toolchain = context.getToolchain(llvm::errs()); - if (!toolchain) + if (!toolchain) { return failure(); + } if (failed(mlir::pto::emitFatobjCCE( jobResult.textOutput, fatobjPath, *toolchain, context.getTempFiles(), llvm::errs()))) { @@ -922,8 +973,9 @@ class VPTOBackendChildJob final : public BackendChildJob { } std::string fatobjPath; - if (failed(context.createTempPath("ptoas-vpto-fatobj", ".o", fatobjPath))) + if (failed(context.createTempPath("ptoas-vpto-fatobj", ".o", fatobjPath))) { return failure(); + } if (failed(emitVPTOLLVMFatobj(jobResult, context, moduleId, fatobjPath))) { dumpFailedMixedChildCompileUnit("vpto", summary, op); @@ -954,12 +1006,14 @@ class FatobjLinkJob { } std::string stderrPath; - if (failed(context.createTempPath("ptoas-fatobj", ".log", stderrPath))) + if (failed(context.createTempPath("ptoas-fatobj", ".log", stderrPath))) { return failure(); + } const mlir::pto::CANNToolchain *toolchain = context.getToolchain(llvm::errs()); - if (!toolchain) + if (!toolchain) { return failure(); + } return mlir::pto::linkFatobjs(fatobjPaths, context.getOutputPath(), *toolchain, stderrPath, llvm::errs()); } @@ -979,8 +1033,9 @@ LogicalResult EmitCBackendJob::run(PTOASContext &context) { isBackendPartitionedContainer(op)) { FailureOr> jobModuleOr = buildBackendChildCompileUnit(op, children.front()); - if (failed(jobModuleOr)) + if (failed(jobModuleOr)) { return failure(); + } singleChildJobModule = std::move(*jobModuleOr); singleChildJobModule.get()->setAttr( "pto.backend", @@ -1015,8 +1070,9 @@ LogicalResult VPTOBackendJob::run(PTOASContext &context) { isBackendPartitionedContainer(op)) { FailureOr> jobModuleOr = buildBackendChildCompileUnit(op, children.front()); - if (failed(jobModuleOr)) + if (failed(jobModuleOr)) { return failure(); + } singleChildJobModule = std::move(*jobModuleOr); singleChildJobModule.get()->setAttr( "pto.backend", @@ -1030,8 +1086,9 @@ LogicalResult VPTOBackendJob::run(PTOASContext &context) { *compileUnit, context, mlir::pto::PTOBackend::VPTO, result, emitHostStub) != 0) return failure(); - if (result.kind == mlir::pto::PTOASCompileResultKind::Text) + if (result.kind == mlir::pto::PTOASCompileResultKind::Text) { return success(); + } if (result.kind != mlir::pto::PTOASCompileResultKind::VPTOObject) { llvm::errs() << "Error: VPTO backend job produced non-VPTO output.\n"; return failure(); @@ -1057,13 +1114,15 @@ static LogicalResult emitVPTOLLVMFatobj( const mlir::pto::PTOASCompileResult &jobResult, PTOASContext &context, llvm::StringRef moduleId, llvm::StringRef outputPath) { llvm::StringRef stubSource = kEmptyHostStubSource; - if (!jobResult.vptoStubSource.empty()) + if (!jobResult.vptoStubSource.empty()) { stubSource = jobResult.vptoStubSource; + } const mlir::pto::CANNToolchain *toolchain = context.getToolchain(llvm::errs()); - if (!toolchain) + if (!toolchain) { return failure(); + } if (failed(mlir::pto::emitFatobjLLVM( jobResult.vptoCubeModule.module.get(), jobResult.vptoVectorModule.module.get(), stubSource, @@ -1081,13 +1140,15 @@ static LogicalResult collectChildJobs( SmallVector children(module.getOps()); for (ModuleOp child : children) { std::optional childBackend; - if (failed(parseDriverBackendAttr(child.getOperation(), childBackend))) + if (failed(parseDriverBackendAttr(child.getOperation(), childBackend))) { return failure(); + } FailureOr> jobModuleOr = buildBackendChildCompileUnit(module, child); - if (failed(jobModuleOr)) + if (failed(jobModuleOr)) { return failure(); + } OwningOpRef jobModule = std::move(*jobModuleOr); if (llvm::sys::Process::GetEnv("PTOAS_DEBUG_CHILD_UNIT")) { llvm::errs() << "// ----- child compile unit ----- //\n"; @@ -1143,8 +1204,9 @@ static LogicalResult resolveSingleBackend( std::optional firstChildBackend; for (ModuleOp child : children) { std::optional childBackend; - if (failed(parseDriverBackendAttr(child.getOperation(), childBackend))) + if (failed(parseDriverBackendAttr(child.getOperation(), childBackend))) { return failure(); + } mlir::pto::PTOBackend effectiveChildBackend = childBackend.value_or(defaultBackend); @@ -1152,14 +1214,17 @@ static LogicalResult resolveSingleBackend( firstChildBackend = effectiveChildBackend; continue; } - if (*firstChildBackend != effectiveChildBackend) + if (*firstChildBackend != effectiveChildBackend) { return success(); + } } - if (firstChildBackend) + if (firstChildBackend) { singleBackend = *firstChildBackend; - else + } + else { singleBackend = defaultBackend; + } return success(); } @@ -1176,8 +1241,9 @@ static LogicalResult buildBackendInfo(ModuleOp module, bool cliBackendSpecified, std::optional moduleBackend; if (!cliBackendSpecified) { - if (failed(parseDriverBackendAttr(module.getOperation(), moduleBackend))) + if (failed(parseDriverBackendAttr(module.getOperation(), moduleBackend))) { return failure(); + } } if (failed(resolveSingleBackend(cliBackendSpecified, moduleBackend, @@ -1234,13 +1300,15 @@ static LogicalResult runPTOASJobs(OwningOpRef &module, result.kind = mlir::pto::PTOASCompileResultKind::MixedObject; for (size_t i = 0, e = backendJobs.size(); i < e; ++i) { - if (failed(backendJobs[i]->run(context))) + if (failed(backendJobs[i]->run(context))) { return failure(); + } } FatobjLinkJob linkJob(fatobjPaths); - if (failed(linkJob.run(context))) + if (failed(linkJob.run(context))) { return failure(); + } return success(); } @@ -1277,8 +1345,9 @@ static int runPTOASDriver(int argc, char **argv, MLIRContext *borrowedContext = nullptr) { DialectRegistry registry; mlir::pto::registerPTOASDialects(registry); - if (borrowedContext) + if (borrowedContext) { borrowedContext->appendDialectRegistry(registry); + } mlir::pto::registerPTOASPassesAndCLOptions(); llvm::cl::SetVersionPrinter(printPTOASVersion); @@ -1311,31 +1380,37 @@ static int runPTOASDriver(int argc, char **argv, context->initializeMLIRContext(); std::unique_ptr inputBuffer = readInputBuffer(); - if (!inputBuffer) + if (!inputBuffer) { return 1; + } std::string arch; OwningOpRef module = loadInputModule( std::move(inputBuffer), context->getMLIRContext(), cliArchSpecified, arch); - if (!module) + if (!module) { return 1; + } context->setArch(std::move(arch)); mlir::pto::BackendInfo backendInfo; - if (failed(buildBackendInfo(module.get(), cliBackendSpecified, backendInfo))) + if (failed(buildBackendInfo(module.get(), cliBackendSpecified, backendInfo))) { return 1; + } context->setBackendInfo(std::move(backendInfo)); (void)context->initializeEnvironment( context->getBackendInfo().requiresToolchain, llvm::errs()); mlir::pto::PTOASCompileResult result; - if (failed(runPTOASJobs(module, *context, result))) + if (failed(runPTOASJobs(module, *context, result))) { return 1; + } - if (result.kind == mlir::pto::PTOASCompileResultKind::Text) + if (result.kind == mlir::pto::PTOASCompileResultKind::Text) { return failed(writeTextOutput(result.textOutput, context->getOutputPath())); - if (result.kind == mlir::pto::PTOASCompileResultKind::MixedObject) + } + if (result.kind == mlir::pto::PTOASCompileResultKind::MixedObject) { return 0; + } llvm::errs() << "Error: unsupported ptoas compile result.\n"; return 1; diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 4c89890f09..45be409cef 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -77,7 +77,7 @@ #include #include #include -#include +#include #include extern "C" { @@ -116,15 +116,17 @@ struct ApplySIMTEntryNoInlinePass final void runOnOperation() final { for (func::FuncOp func : getOperation().getOps()) - if (func->hasAttr(pto::kPTOSimtEntryAttrName)) + if (func->hasAttr(pto::kPTOSimtEntryAttrName)) { func.setNoInline(true); + } } }; static std::string normalizeArch(llvm::StringRef arch) { std::string normalized = arch.str(); - for (char &c : normalized) + for (char &c : normalized) { c = static_cast(std::tolower(static_cast(c))); + } return normalized; } @@ -140,37 +142,44 @@ static bool isSupportedPTOASTargetArch(llvm::StringRef arch) { static std::optional getModuleTargetArchAttr(ModuleOp module) { auto attr = module->getAttrOfType("pto.target_arch"); - if (!attr) + if (!attr) { return std::nullopt; + } std::string arch = normalizeArch(attr.getValue()); - if (!isSupportedPTOASTargetArch(arch)) + if (!isSupportedPTOASTargetArch(arch)) { return std::nullopt; + } return arch; } static std::string resolveEffectiveTargetArch(ModuleOp module, llvm::StringRef fallbackArch) { - if (std::optional arch = getModuleTargetArchAttr(module)) + if (std::optional arch = getModuleTargetArchAttr(module)) { return *arch; + } std::optional childArch; for (ModuleOp child : module.getOps()) { std::optional arch = getModuleTargetArchAttr(child); - if (!arch) + if (!arch) { continue; + } if (!childArch) { childArch = std::move(arch); continue; } - if (*childArch != *arch) + if (*childArch != *arch) { return normalizeArch(fallbackArch); + } } - if (childArch) + if (childArch) { return *childArch; + } std::string fallback = normalizeArch(fallbackArch); - if (!isSupportedPTOASTargetArch(fallback)) + if (!isSupportedPTOASTargetArch(fallback)) { return "a3"; + } return fallback; } @@ -231,8 +240,9 @@ void mlir::pto::loadPTOASDialects(MLIRContext &context) { static LogicalResult applyConfiguredPassManagerCLOptions( PassManager &pm, llvm::StringRef pipelineName, llvm::raw_ostream &diagOS = llvm::errs()) { - if (succeeded(mlir::applyPassManagerCLOptions(pm))) + if (succeeded(mlir::applyPassManagerCLOptions(pm))) { return success(); + } diagOS << "Error: failed to apply MLIR pass manager command-line options for " << pipelineName << ".\n"; return failure(); @@ -254,8 +264,9 @@ static LogicalResult reorderEmitCFunctions(ModuleOp module) { llvm::DenseMap indegree; llvm::DenseMap> outgoing; - for (auto func : definitions) + for (auto func : definitions) { indegree[func.getOperation()] = 0; + } for (auto caller : definitions) { Operation *callerOp = caller.getOperation(); @@ -263,18 +274,21 @@ static LogicalResult reorderEmitCFunctions(ModuleOp module) { bool hasCycle = false; caller.walk([&](emitc::CallOp call) { auto calleeAttr = call.getCalleeAttr(); - if (!calleeAttr) + if (!calleeAttr) { return; + } auto it = definitionsByName.find(calleeAttr.getLeafReference()); - if (it == definitionsByName.end()) + if (it == definitionsByName.end()) { return; + } Operation *calleeOp = it->second.getOperation(); if (calleeOp == callerOp) { hasCycle = true; return; } - if (!seenCallees.insert(calleeOp).second) + if (!seenCallees.insert(calleeOp).second) { return; + } outgoing[calleeOp].push_back(callerOp); ++indegree[callerOp]; }); @@ -287,8 +301,9 @@ static LogicalResult reorderEmitCFunctions(ModuleOp module) { SmallVector ready; for (auto func : definitions) { - if (indegree[func.getOperation()] == 0) + if (indegree[func.getOperation()] == 0) { ready.push_back(func.getOperation()); + } } SmallVector sortedDefinitions; @@ -300,8 +315,9 @@ static LogicalResult reorderEmitCFunctions(ModuleOp module) { for (Operation *user : outgoing[next]) { unsigned &userIndegree = indegree[user]; - if (--userIndegree == 0) + if (--userIndegree == 0) { ready.push_back(user); + } } } @@ -310,8 +326,9 @@ static LogicalResult reorderEmitCFunctions(ModuleOp module) { << "cyclic function call graph is not supported for EmitC C++ emission"; } - if (declarations.empty() && definitions.size() <= 1) + if (declarations.empty() && definitions.size() <= 1) { return success(); + } SmallVector desiredOrder; desiredOrder.append(declarations.begin(), declarations.end()); @@ -325,14 +342,16 @@ static LogicalResult reorderEmitCFunctions(ModuleOp module) { break; } } - if (!anchor) + if (!anchor) { return success(); + } auto advanceAnchor = [&]() { while (anchor) { anchor = anchor->getNextNode(); - if (!anchor || isa(anchor)) + if (!anchor || isa(anchor)) { return; + } } }; @@ -341,10 +360,12 @@ static LogicalResult reorderEmitCFunctions(ModuleOp module) { advanceAnchor(); continue; } - if (anchor) + if (anchor) { func->moveBefore(anchor); - else + } + else { func->moveBefore(&body, body.end()); + } } return success(); @@ -552,8 +573,9 @@ static PTOBuildLevel defaultBuildLevel() { static bool parseBuildLevel(llvm::StringRef levelStr, PTOBuildLevel &out) { std::string s = levelStr.str(); - for (char &c : s) + for (char &c : s) { c = static_cast(std::tolower(static_cast(c))); + } if (s == "level1") { out = PTOBuildLevel::Level1; return true; @@ -578,9 +600,9 @@ static ReserveBufferMemSpec getReserveBufferMemSpec(PTOArch arch, AddressSpace space) { switch (space) { case AddressSpace::VEC: - return {arch == PTOArch::A5 ? 253952ull : 196608ull, 256}; + return {arch == PTOArch::A5 ? 253952uLL : 196608uLL, 256}; case AddressSpace::MAT: - return {524288ull, 256}; + return {524288uLL, 256}; case AddressSpace::LEFT: case AddressSpace::RIGHT: case AddressSpace::ACC: @@ -596,12 +618,14 @@ static ReserveBufferMemSpec getReserveBufferMemSpec(PTOArch arch, static LogicalResult validateReserveBufferBase(pto::ReserveBufferOp op, PTOArch arch) { auto baseAttr = op.getBaseAttr(); - if (!baseAttr) + if (!baseAttr) { return op.emitError("expects explicit 'base'"); + } int64_t signedBase = baseAttr.getInt(); - if (signedBase < 0) + if (signedBase < 0) { return op.emitError("expects 'base' to be non-negative when present"); + } ReserveBufferMemSpec spec = getReserveBufferMemSpec(arch, op.getLocation().getAddressSpace()); @@ -638,8 +662,9 @@ static bool validateReserveBufferLevelRules(ModuleOp module, return; } - if (op.getBaseAttr()) + if (op.getBaseAttr()) { (void)validateReserveBufferBase(op, arch); + } op.emitError("pto.reserve_buffer with explicit 'base' (auto = false) is " "not supported when --pto-level=level1 or level2; use " "--pto-level=level3 or set auto = true"); @@ -654,8 +679,9 @@ static bool validateReserveBufferLevelRules(ModuleOp module, return; } - if (mlir::failed(validateReserveBufferBase(op, arch))) + if (mlir::failed(validateReserveBufferBase(op, arch))) { failed = true; + } }); return !failed; } @@ -667,8 +693,9 @@ static constexpr llvm::StringLiteral kAutoSyncTailPolicyMte3ToSEvent0 = static bool parseAutoSyncTailHint(llvm::StringRef hintStr, std::string &normalized) { std::string s = hintStr.str(); - for (char &c : s) + for (char &c : s) { c = static_cast(std::tolower(static_cast(c))); + } if (s == "barrier-all" || s == "barrier_all" || s == "default") { normalized = kAutoSyncTailPolicyBarrierAll.str(); return true; @@ -684,8 +711,9 @@ static bool parseAutoSyncTailHint(llvm::StringRef hintStr, std::string &normaliz static LogicalResult emitSharedPreBackendSeamIR(ModuleOp module, llvm::StringRef outputPath) { - if (outputPath.empty()) + if (outputPath.empty()) { return success(); + } if (outputPath == "-") { module->print(llvm::outs()); @@ -716,8 +744,9 @@ static void printSharedPreBackendSeamIR(ModuleOp module) { static bool hasUnexpandedTileOps(ModuleOp module) { bool found = false; module.walk([&](Operation *op) { - if (found) + if (found) { return; + } if (isa(op)) { found = true; return; @@ -749,18 +778,22 @@ static bool isCppIdentifierChar(char c) { } static std::optional getTextualNameFromSMRange(llvm::SMRange range) { - if (!range.Start.isValid() || !range.End.isValid()) + if (!range.Start.isValid() || !range.End.isValid()) { return std::nullopt; + } const char *begin = range.Start.getPointer(); const char *end = range.End.getPointer(); - if (!begin || !end || end < begin) + if (!begin || !end || end < begin) { return std::nullopt; + } llvm::StringRef name(begin, static_cast(end - begin)); - if (name.empty()) + if (name.empty()) { return std::nullopt; + } name = name.trim(); - if (name.consume_front("%") && name.empty()) + if (name.consume_front("%") && name.empty()) { return std::nullopt; + } return name.str(); } @@ -768,26 +801,30 @@ static SmallVector expandTextualResultGroupHints(const AsmParserState::OperationDefinition &opDef, unsigned groupIndex) { SmallVector hints; - if (groupIndex >= opDef.resultGroups.size()) + if (groupIndex >= opDef.resultGroups.size()) { return hints; + } const auto &group = opDef.resultGroups[groupIndex]; std::optional baseName = getTextualNameFromSMRange(group.definition.loc); - if (!baseName) + if (!baseName) { return hints; + } unsigned resultStart = group.startIndex; unsigned resultEnd = groupIndex + 1 == opDef.resultGroups.size() ? opDef.op->getNumResults() : opDef.resultGroups[groupIndex + 1].startIndex; - if (resultStart >= resultEnd) + if (resultStart >= resultEnd) { return hints; + } if (resultEnd - resultStart == 1) { hints.push_back(*baseName); return hints; } - for (unsigned idx = resultStart; idx < resultEnd; ++idx) + for (unsigned idx = resultStart; idx < resultEnd; ++idx) { hints.push_back(*baseName + "#" + std::to_string(idx - resultStart)); + } return hints; } @@ -796,24 +833,29 @@ static std::string sanitizeCppIdentifier(llvm::StringRef name) { sanitized.reserve(name.size() + 4); auto appendUnderscore = [&]() { - if (sanitized.empty() || sanitized.back() != '_') + if (sanitized.empty() || sanitized.back() != '_') { sanitized.push_back('_'); + } }; for (char c : name) { - if (isCppIdentifierChar(c)) + if (isCppIdentifierChar(c)) { sanitized.push_back(c); - else + } + else { appendUnderscore(); + } } - while (!sanitized.empty() && sanitized.back() == '_') + while (!sanitized.empty() && sanitized.back() == '_') { sanitized.pop_back(); + } if (sanitized.empty()) return {}; - if (!isCppIdentifierStart(sanitized.front())) + if (!isCppIdentifierStart(sanitized.front())) { sanitized.insert(sanitized.begin(), '_'); + } return sanitized; } @@ -821,8 +863,9 @@ static void appendLocationNameHints(Location loc, SmallVectorImpl &hints) { if (auto nameLoc = dyn_cast(loc)) { std::string sanitized = sanitizeCppIdentifier(nameLoc.getName().getValue()); - if (!sanitized.empty()) + if (!sanitized.empty()) { hints.push_back(std::move(sanitized)); + } return; } @@ -830,21 +873,25 @@ static void appendLocationNameHints(Location loc, if (Attribute metadata = fusedLoc.getMetadata()) { if (auto strAttr = dyn_cast(metadata)) { std::string sanitized = sanitizeCppIdentifier(strAttr.getValue()); - if (!sanitized.empty()) + if (!sanitized.empty()) { hints.push_back(std::move(sanitized)); + } return; } if (auto arrayAttr = dyn_cast(metadata)) { for (Attribute attr : arrayAttr) { auto strAttr = dyn_cast(attr); - if (!strAttr) + if (!strAttr) { continue; + } std::string sanitized = sanitizeCppIdentifier(strAttr.getValue()); - if (!sanitized.empty()) + if (!sanitized.empty()) { hints.push_back(std::move(sanitized)); + } } - if (!hints.empty()) + if (!hints.empty()) { return; + } } } @@ -856,8 +903,9 @@ static void appendLocationNameHints(Location loc, if (auto callSiteLoc = dyn_cast(loc)) { appendLocationNameHints(callSiteLoc.getCallee(), hints); - if (hints.empty()) + if (hints.empty()) { appendLocationNameHints(callSiteLoc.getCaller(), hints); + } } } @@ -876,8 +924,9 @@ static void appendRawLocationProvenance(Location loc, SmallVectorImpl &hints) { if (auto nameLoc = dyn_cast(loc)) { std::string raw = nameLoc.getName().getValue().str(); - if (!raw.empty()) + if (!raw.empty()) { hints.push_back(std::move(raw)); + } return; } @@ -885,21 +934,25 @@ static void appendRawLocationProvenance(Location loc, if (Attribute metadata = fusedLoc.getMetadata()) { if (auto strAttr = dyn_cast(metadata)) { std::string raw = strAttr.getValue().str(); - if (!raw.empty()) + if (!raw.empty()) { hints.push_back(std::move(raw)); + } return; } if (auto arrayAttr = dyn_cast(metadata)) { for (Attribute attr : arrayAttr) { auto strAttr = dyn_cast(attr); - if (!strAttr) + if (!strAttr) { continue; + } std::string raw = strAttr.getValue().str(); - if (!raw.empty()) + if (!raw.empty()) { hints.push_back(std::move(raw)); + } } - if (!hints.empty()) + if (!hints.empty()) { return; + } } } @@ -911,8 +964,9 @@ static void appendRawLocationProvenance(Location loc, if (auto callSiteLoc = dyn_cast(loc)) { appendRawLocationProvenance(callSiteLoc.getCallee(), hints); - if (hints.empty()) + if (hints.empty()) { appendRawLocationProvenance(callSiteLoc.getCaller(), hints); + } } } @@ -921,25 +975,30 @@ static void appendRawLocationProvenance(Location loc, // but without sanitization. static SmallVector getRawResultProvenance(Operation *op) { SmallVector hints; - if (!op || op->getNumResults() == 0) + if (!op || op->getNumResults() == 0) { return hints; + } appendRawLocationProvenance(op->getLoc(), hints); - if (hints.empty()) + if (hints.empty()) { return hints; + } hints.erase(std::remove_if(hints.begin(), hints.end(), [](const std::string &name) { return name.empty(); }), hints.end()); - if (hints.empty()) + if (hints.empty()) { return hints; + } if (op->getNumResults() == 1) { - if (hints.size() > 1) + if (hints.size() > 1) { hints.resize(1); + } return hints; } - if (hints.size() > op->getNumResults()) + if (hints.size() > op->getNumResults()) { hints.resize(op->getNumResults()); + } return hints; } @@ -956,8 +1015,9 @@ static SmallVector getRawLocationProvenance(Location loc) { static Location getIndexedRawProvenanceLoc(Location fallbackLoc, unsigned index) { SmallVector hints = getRawLocationProvenance(fallbackLoc); - if (index >= hints.size()) + if (index >= hints.size()) { return fallbackLoc; + } return NameLoc::get(StringAttr::get(fallbackLoc.getContext(), hints[index]), fallbackLoc); } @@ -968,20 +1028,24 @@ static Location attachLocationNameHints(Location baseLoc, SmallVector attrs; attrs.reserve(hints.size()); for (llvm::StringRef hint : hints) { - if (!hint.empty()) + if (!hint.empty()) { attrs.push_back(StringAttr::get(context, hint)); + } } - if (attrs.empty()) + if (attrs.empty()) { return baseLoc; - if (attrs.size() == 1) + } + if (attrs.size() == 1) { return NameLoc::get(cast(attrs.front()), baseLoc); + } return FusedLoc::get(ArrayRef{baseLoc}, ArrayAttr::get(context, attrs), context); } static void applyValueNameHints(Value value, llvm::ArrayRef hints) { - if (!value || hints.empty() || hasLocationNameHints(value.getLoc())) + if (!value || hints.empty() || hasLocationNameHints(value.getLoc())) { return; + } value.setLoc(attachLocationNameHints(value.getLoc(), hints, value.getContext())); } @@ -996,8 +1060,9 @@ static void applyOperationResultNameHints(Operation *op, for (size_t i = 0, e = std::min(op->getNumResults(), hints.size()); i < e; ++i) limitedHints.push_back(hints[i]); - if (limitedHints.empty()) + if (limitedHints.empty()) { return; + } op->setLoc(attachLocationNameHints(op->getLoc(), limitedHints, op->getContext())); } @@ -1007,8 +1072,9 @@ static void splitDerivedSingleResultProvenanceLocsInRegion(Region ®ion); static void splitDerivedSingleResultProvenanceLocsInBlock(Block &block) { SmallVector ops; ops.reserve(block.getOperations().size()); - for (Operation &op : block) + for (Operation &op : block) { ops.push_back(&op); + } for (size_t i = 0; i < ops.size();) { Operation *op = ops[i]; @@ -1032,52 +1098,62 @@ static void splitDerivedSingleResultProvenanceLocsInBlock(Block &block) { size_t runSize = runEnd - i; if (runSize == hints.size()) { Location sharedLoc = op->getLoc(); - for (size_t j = 0; j < runSize; ++j) + for (size_t j = 0; j < runSize; ++j) { ops[i + j]->setLoc(getIndexedRawProvenanceLoc(sharedLoc, j)); + } } i = runEnd; } for (Operation &op : block) { - for (Region ®ion : op.getRegions()) + for (Region ®ion : op.getRegions()) { splitDerivedSingleResultProvenanceLocsInRegion(region); + } } } static void splitDerivedSingleResultProvenanceLocsInRegion(Region ®ion) { - for (Block &block : region) + for (Block &block : region) { splitDerivedSingleResultProvenanceLocsInBlock(block); + } } static void splitDerivedSingleResultProvenanceLocs(Operation *root) { - if (!root) + if (!root) { return; - for (Region ®ion : root->getRegions()) + } + for (Region ®ion : root->getRegions()) { splitDerivedSingleResultProvenanceLocsInRegion(region); + } } static void narrowUnusedMultiResultProvenanceLocs(Operation *root) { - if (!root) + if (!root) { return; + } root->walk([&](Operation *op) { - if (op->getNumResults() <= 1) + if (op->getNumResults() <= 1) { return; + } SmallVector hints = getRawLocationProvenance(op->getLoc()); - if (hints.size() != op->getNumResults()) + if (hints.size() != op->getNumResults()) { return; + } SmallVector liveHints; liveHints.reserve(hints.size()); for (auto [index, result] : llvm::enumerate(op->getResults())) { - if (!result.use_empty()) + if (!result.use_empty()) { liveHints.push_back(hints[index]); + } } - if (liveHints.empty() || liveHints.size() == hints.size()) + if (liveHints.empty() || liveHints.size() == hints.size()) { return; + } op->setLoc(attachLocationNameHints(op->getLoc(), liveHints, op->getContext())); @@ -1206,37 +1282,44 @@ static void collectNonEntryBlocksInSourceOrder( for (Region ®ion : op->getRegions()) { bool isEntryBlock = true; for (Block &block : region) { - if (!isEntryBlock && block.getNumArguments() != 0) + if (!isEntryBlock && block.getNumArguments() != 0) { blocks.push_back(&block); + } isEntryBlock = false; - for (Operation &nestedOp : block) + for (Operation &nestedOp : block) { collectNonEntryBlocksInSourceOrder(&nestedOp, blocks); + } } } } void mlir::pto::applyTextualNameHintsToModule(ModuleOp module, const AsmParserState &parserState) { - if (!module) + if (!module) { return; + } for (const AsmParserState::BlockDefinition &blockDef : parserState.getBlockDefs()) { - if (!blockDef.block) + if (!blockDef.block) { continue; + } for (auto [argIndex, argDef] : llvm::enumerate(blockDef.arguments)) { - if (argIndex >= blockDef.block->getNumArguments()) + if (argIndex >= blockDef.block->getNumArguments()) { break; + } std::optional hint = getTextualNameFromSMRange(argDef.loc); - if (!hint) + if (!hint) { continue; + } applyValueNameHints(blockDef.block->getArgument(argIndex), llvm::ArrayRef{*hint}); } } for (const AsmParserState::OperationDefinition &opDef : parserState.getOpDefs()) { - if (!opDef.op || opDef.op->getNumResults() == 0) + if (!opDef.op || opDef.op->getNumResults() == 0) { continue; + } SmallVector hints; hints.reserve(opDef.op->getNumResults()); @@ -1246,8 +1329,9 @@ void mlir::pto::applyTextualNameHintsToModule(ModuleOp module, expandTextualResultGroupHints(opDef, groupIndex); hints.append(groupHints.begin(), groupHints.end()); } - if (hints.empty()) + if (hints.empty()) { continue; + } applyOperationResultNameHints(opDef.op, hints); } } @@ -1257,8 +1341,9 @@ static FunctionBlockArgHintMap collectFunctionBlockArgNameHints(ModuleOp module) for (func::FuncOp func : module.getOps()) { SmallVector nonEntryBlocks; collectNonEntryBlocksInSourceOrder(func.getOperation(), nonEntryBlocks); - if (nonEntryBlocks.empty()) + if (nonEntryBlocks.empty()) { continue; + } SmallVector, 4> blockHints; blockHints.reserve(nonEntryBlocks.size()); @@ -1273,12 +1358,14 @@ static FunctionBlockArgHintMap collectFunctionBlockArgNameHints(ModuleOp module) } argHints.push_back(std::move(hints.front())); } - if (hasAllHints) + if (hasAllHints) { blockHints.push_back(std::move(argHints)); + } } - if (!blockHints.empty()) + if (!blockHints.empty()) { hintsByFunction[func.getSymNameAttr()] = std::move(blockHints); + } } return hintsByFunction; } @@ -1287,13 +1374,15 @@ static void applyFunctionBlockArgNameHintsToEmitC( ModuleOp module, const FunctionBlockArgHintMap &blockArgHints) { for (emitc::FuncOp func : module.getOps()) { auto it = blockArgHints.find(func.getSymNameAttr()); - if (it == blockArgHints.end() || it->second.empty()) + if (it == blockArgHints.end() || it->second.empty()) { continue; + } SmallVector nonEntryBlocks; collectNonEntryBlocksInSourceOrder(func.getOperation(), nonEntryBlocks); - if (nonEntryBlocks.size() != it->second.size()) + if (nonEntryBlocks.size() != it->second.size()) { continue; + } bool shapeMatches = true; for (auto [blockIndex, block] : llvm::enumerate(nonEntryBlocks)) { @@ -1302,8 +1391,9 @@ static void applyFunctionBlockArgNameHintsToEmitC( break; } } - if (!shapeMatches) + if (!shapeMatches) { continue; + } for (auto [blockIndex, block] : llvm::enumerate(nonEntryBlocks)) { const auto &argHints = it->second[blockIndex]; @@ -1315,11 +1405,13 @@ static void applyFunctionBlockArgNameHintsToEmitC( static SmallVector getValueNameHints(Value value) { SmallVector hints; - if (!value) + if (!value) { return hints; + } appendLocationNameHints(value.getLoc(), hints); - if (hints.size() > 1) + if (hints.size() > 1) { hints.resize(1); + } return hints; } @@ -1365,8 +1457,9 @@ collectExpressionProvenance(emitc::ExpressionOp expr) { SmallVector provenance; auto appendUnique = [&](llvm::ArrayRef names) { for (const std::string &name : names) { - if (name.empty()) + if (name.empty()) { continue; + } if (std::find(provenance.begin(), provenance.end(), name) != provenance.end()) continue; @@ -1375,10 +1468,12 @@ collectExpressionProvenance(emitc::ExpressionOp expr) { }; expr.walk([&](Operation *nested) { - if (nested == expr.getOperation()) + if (nested == expr.getOperation()) { return WalkResult::advance(); - if (nested->getNumResults() == 0 || isa(nested)) + } + if (nested->getNumResults() == 0 || isa(nested)) { return WalkResult::advance(); + } appendUnique(getRawResultProvenance(nested)); return WalkResult::advance(); }); @@ -1394,26 +1489,30 @@ static void annotateEmitCProvenanceHints(ModuleOp module) { llvm::SmallVector opsToAnnotate; module.walk([&](Operation *op) { - if (op->getNumResults() == 0 || isa(op)) + if (op->getNumResults() == 0 || isa(op)) { return WalkResult::advance(); + } if (auto expr = dyn_cast(op)) { SmallVector provenance = collectExpressionProvenance(expr); - if (provenance.empty()) + if (provenance.empty()) { return WalkResult::skip(); + } opsToAnnotate.push_back( ProvenanceMarker{op, SmallVector(provenance)}); return WalkResult::skip(); } - if (op->getParentOfType()) + if (op->getParentOfType()) { return WalkResult::advance(); + } // Only carry raw provenance into the C++ post-pass. Semantic renaming is // intentionally deferred until naming can happen inside the emitter's own // symbol table instead of via post-hoc C++ text rewriting. SmallVector provenance = getRawResultProvenance(op); - if (provenance.empty()) + if (provenance.empty()) { return WalkResult::advance(); + } opsToAnnotate.push_back(ProvenanceMarker{ op, SmallVector(provenance.begin(), provenance.end())}); return WalkResult::advance(); @@ -1479,8 +1578,9 @@ static bool parseMarkerArgs(llvm::StringRef argsRef, continue; } if (c == ')') { - if (parenDepth > 0) + if (parenDepth > 0) { --parenDepth; + } continue; } if (c == ',' && parenDepth == 0) { @@ -1488,8 +1588,9 @@ static bool parseMarkerArgs(llvm::StringRef argsRef, partBegin = i + 1; } } - if (partBegin > argsRef.size()) + if (partBegin > argsRef.size()) { return false; + } args.push_back(argsRef.drop_front(partBegin).trim()); return true; } @@ -1499,8 +1600,9 @@ findNextMarkerCall(const std::string &cpp, llvm::StringRef marker, size_t searchPos) { ParsedMarkerCall call; call.markerPos = cpp.find(marker.str(), searchPos); - if (call.markerPos == std::string::npos) + if (call.markerPos == std::string::npos) { return std::nullopt; + } size_t lparenPos = call.markerPos + marker.size(); if (lparenPos >= cpp.size() || cpp[lparenPos] != '(') @@ -1514,20 +1616,23 @@ findNextMarkerCall(const std::string &cpp, llvm::StringRef marker, ++parenDepth; continue; } - if (c != ')') + if (c != ')') { continue; + } if (parenDepth == 0) { call.rparenPos = i; break; } --parenDepth; } - if (call.rparenPos == std::string::npos) + if (call.rparenPos == std::string::npos) { return call; + } llvm::StringRef argsRef(cpp.data() + argsBegin, call.rparenPos - argsBegin); - if (!parseMarkerArgs(argsRef, call.args)) + if (!parseMarkerArgs(argsRef, call.args)) { call.args.clear(); + } return call; } @@ -1562,8 +1667,9 @@ static bool rewriteMarkerCallToMember(std::string &cpp, llvm::StringRef marker, unsigned expectedNumArgs) { return rewriteMarkerCalls( cpp, marker, [&](const ParsedMarkerCall &call) -> std::optional { - if (call.args.size() != expectedNumArgs) + if (call.args.size() != expectedNumArgs) { return std::nullopt; + } std::string replacement; replacement.reserve(marker.size() + kMarkerCallReserveExtra); @@ -1571,8 +1677,9 @@ static bool rewriteMarkerCallToMember(std::string &cpp, llvm::StringRef marker, replacement.push_back('.'); replacement.append(memberName.str()); replacement.push_back('('); - if (expectedNumArgs >= kMarkerRewriteMinArgCount) + if (expectedNumArgs >= kMarkerRewriteMinArgCount) { replacement.append(call.args[1].str()); + } if (expectedNumArgs == kMarkerRewriteTernaryArgCount) { replacement.append(", "); replacement.append(call.args[2].str()); @@ -1600,10 +1707,12 @@ static bool rewriteMarkerCallToField(std::string &cpp, llvm::StringRef marker, size_t expectedNumArgs) { return rewriteMarkerCalls( cpp, marker, [&](const ParsedMarkerCall &call) -> std::optional { - if (call.args.size() != expectedNumArgs) + if (call.args.size() != expectedNumArgs) { return std::nullopt; - if (call.args.empty()) + } + if (call.args.empty()) { return std::nullopt; + } std::string replacement; replacement.reserve(call.args.front().size() + fieldName.size() + 1); replacement.append(call.args.front().str()); @@ -1650,21 +1759,25 @@ static void dropEmptyEmitCExpressions(Operation *rootOp) { toErase; rootOp->walk([&](emitc::ExpressionOp expr) { Block *body = expr.getBody(); - if (!body) + if (!body) { return; + } auto yield = dyn_cast(body->getTerminator()); - if (!yield || yield.getNumOperands() != 1) + if (!yield || yield.getNumOperands() != 1) { return; + } Value yielded = yield.getOperand(0); Operation *defOp = yielded.getDefiningOp(); bool yieldedFromOutside = !defOp || defOp->getBlock() != body; - if (!yieldedFromOutside && expr.getRootOp()) + if (!yieldedFromOutside && expr.getRootOp()) { return; + } expr.getResult().replaceAllUsesWith(yielded); toErase.push_back(expr); }); - for (emitc::ExpressionOp expr : llvm::reverse(toErase)) + for (emitc::ExpressionOp expr : llvm::reverse(toErase)) { expr.erase(); + } } static void appendEmitCIntegerAttrLiteral(std::string &storage, @@ -1698,8 +1811,9 @@ static std::string getEmitCIntegerAttrLiteral(IntegerAttr attr) { static std::optional getEmitCDenseIntElementsAttrLiteral(DenseIntElementsAttr attr) { auto tensorTy = dyn_cast(attr.getType()); - if (!tensorTy) + if (!tensorTy) { return std::nullopt; + } Type elementType = tensorTy.getElementType(); bool isUnsigned = false; @@ -1713,8 +1827,9 @@ getEmitCDenseIntElementsAttrLiteral(DenseIntElementsAttr attr) { literal.push_back('{'); bool first = true; for (const APInt &value : attr) { - if (!first) + if (!first) { literal.append(", "); + } first = false; appendEmitCIntegerAttrLiteral(literal, value, isUnsigned); } @@ -1724,8 +1839,9 @@ getEmitCDenseIntElementsAttrLiteral(DenseIntElementsAttr attr) { static Attribute normalizeEmitCPrintedAttrForCppEmission(MLIRContext *ctx, Attribute attr) { - if (auto intAttr = dyn_cast(attr)) + if (auto intAttr = dyn_cast(attr)) { return emitc::OpaqueAttr::get(ctx, getEmitCIntegerAttrLiteral(intAttr)); + } if (auto denseAttr = dyn_cast(attr)) { if (std::optional literal = @@ -1743,8 +1859,9 @@ static Attribute normalizeEmitCPrintedAttrForCppEmission(MLIRContext *ctx, changed |= normalizedElement != element; normalized.push_back(normalizedElement); } - if (changed) + if (changed) { return ArrayAttr::get(ctx, normalized); + } } return attr; @@ -1812,8 +1929,9 @@ static void normalizeEmitCIntegerAttrsForCppEmission(Operation *rootOp) { Attribute value = constant.getValue(); Attribute normalized = normalizeEmitCPrintedAttrForCppEmission(ctx, value); - if (normalized != value) + if (normalized != value) { constant.getProperties().setValue(normalized); + } return; } @@ -1821,33 +1939,38 @@ static void normalizeEmitCIntegerAttrsForCppEmission(Operation *rootOp) { Attribute value = variable.getValue(); Attribute normalized = normalizeEmitCPrintedAttrForCppEmission(ctx, value); - if (normalized != value) + if (normalized != value) { variable.getProperties().setValue(normalized); + } return; } if (auto global = dyn_cast(op)) { std::optional initialValue = global.getInitialValue(); - if (!initialValue) + if (!initialValue) { return; + } Attribute normalized = normalizeEmitCPrintedAttrForCppEmission(ctx, *initialValue); - if (normalized != *initialValue) + if (normalized != *initialValue) { global.getProperties().setInitialValue(normalized); + } return; } if (auto call = dyn_cast(op)) { if (std::optional args = call.getArgs()) { ArrayAttr normalized = normalizeEmitCCallArgsForCppEmission(ctx, *args); - if (normalized != *args) + if (normalized != *args) { call.getProperties().setArgs(normalized); + } } if (std::optional templateArgs = call.getTemplateArgs()) { ArrayAttr normalized = normalizeEmitCTemplateArgsForCppEmission(ctx, *templateArgs); - if (normalized != *templateArgs) + if (normalized != *templateArgs) { call.getProperties().setTemplateArgs(normalized); + } } return; } @@ -1856,22 +1979,27 @@ static void normalizeEmitCIntegerAttrsForCppEmission(Operation *rootOp) { static Attribute getDefaultEmitCVariableInitAttr(OpBuilder &builder, Type type) { if (auto intTy = dyn_cast(type)) { - if (intTy.getWidth() == 0) + if (intTy.getWidth() == 0) { return emitc::OpaqueAttr::get(builder.getContext(), "0"); + } return builder.getIntegerAttr(intTy, 0); } - if (isa(type)) + if (isa(type)) { return builder.getIndexAttr(0); - if (auto floatTy = dyn_cast(type)) + } + if (auto floatTy = dyn_cast(type)) { return builder.getFloatAttr(floatTy, 0.0); - if (isa(type)) + } + if (isa(type)) { return emitc::OpaqueAttr::get(builder.getContext(), ""); + } return Attribute{}; } static Type getEmitCVariableStorageType(Type valueType) { - if (isa(valueType)) + if (isa(valueType)) { return valueType; + } return emitc::LValueType::get(valueType); } @@ -1882,8 +2010,9 @@ static Type getEmitCVariableStorageType(Type valueType) { static void materializeControlFlowOperands(Operation *rootOp) { llvm::SmallVector branches; rootOp->walk([&](Operation *op) { - if (isa(op)) + if (isa(op)) { branches.push_back(op); + } }); OpBuilder builder(rootOp->getContext()); @@ -1892,13 +2021,15 @@ static void materializeControlFlowOperands(Operation *rootOp) { for (OpOperand &operand : op->getOpOperands()) { Value value = operand.get(); auto expr = dyn_cast_or_null(value.getDefiningOp()); - if (!expr) + if (!expr) { continue; + } Attribute initAttr = getDefaultEmitCVariableInitAttr(builder, value.getType()); - if (!initAttr) + if (!initAttr) { continue; + } Value tmp = builder .create( @@ -1924,8 +2055,9 @@ static bool rewriteMarkerCallToSubscript(std::string &cpp, llvm::StringRef marke bool isStore) { return rewriteMarkerCalls( cpp, marker, [&](const ParsedMarkerCall &call) -> std::optional { - if (call.args.size() != expectedNumArgs) + if (call.args.size() != expectedNumArgs) { return std::nullopt; + } std::string replacement; replacement.reserve(call.args[0].size() + call.args[1].size() + 8 + (isStore ? call.args[2].size() : 0)); @@ -1966,16 +2098,18 @@ static void rewritePtrScalarMarkers(std::string &cpp) { static std::string getLineIndent(llvm::StringRef line) { size_t firstNonSpace = line.find_first_not_of(" \t"); - if (firstNonSpace == llvm::StringRef::npos) + if (firstNonSpace == llvm::StringRef::npos) { return line.str(); + } return line.take_front(firstNonSpace).str(); } static bool isAICOREFunctionStart(llvm::StringRef trimmed) { if (trimmed.empty() || trimmed.starts_with("#") || trimmed.starts_with("//")) return false; - if (!trimmed.contains("AICORE")) + if (!trimmed.contains("AICORE")) { return false; + } return trimmed.contains("("); } @@ -2007,8 +2141,9 @@ static bool stripScalarGMFlushMarkersFromLine(std::string &line) { size_t searchPos = 0; while (true) { auto call = findNextMarkerCall(line, kMarker, searchPos); - if (!call) + if (!call) { break; + } if (call->rparenPos == std::string::npos) { searchPos = call->markerPos + kMarker.size(); continue; @@ -2023,8 +2158,9 @@ static bool stripScalarGMFlushMarkersFromLine(std::string &line) { while (eraseEnd < line.size() && (line[eraseEnd] == ' ' || line[eraseEnd] == '\t')) ++eraseEnd; - if (eraseEnd < line.size() && line[eraseEnd] == ';') + if (eraseEnd < line.size() && line[eraseEnd] == ';') { ++eraseEnd; + } while (eraseEnd < line.size() && (line[eraseEnd] == ' ' || line[eraseEnd] == '\t')) ++eraseEnd; @@ -2040,8 +2176,9 @@ static bool previousSignificantLineIsTailFlushPoint( llvm::ArrayRef lines, size_t index) { for (size_t i = index; i > 0; --i) { llvm::StringRef prev = llvm::StringRef(lines[i - 1]).trim(); - if (prev.empty()) + if (prev.empty()) { continue; + } return prev.starts_with("#endif // __DAV_") || prev.starts_with("ptoas_auto_sync_tail("); } @@ -2052,8 +2189,9 @@ static bool previousSignificantLineIsExitOrTailFlushPoint( llvm::ArrayRef lines, size_t index) { for (size_t i = index; i > 0; --i) { llvm::StringRef prev = llvm::StringRef(lines[i - 1]).trim(); - if (prev.empty()) + if (prev.empty()) { continue; + } return prev.starts_with("return") || prev.starts_with("#endif // __DAV_") || prev.starts_with("ptoas_auto_sync_tail("); @@ -2082,8 +2220,9 @@ static std::string rewriteScalarGMStoreFlushMarkersInFunction( unchanged.reserve(kRewriteOutputReserveExtra); for (size_t i = 0; i < lines.size(); ++i) { unchanged.append(lines[i]); - if (i + 1 < lines.size() || hasTrailingNewline) + if (i + 1 < lines.size() || hasTrailingNewline) { unchanged.push_back('\n'); + } } return unchanged; } @@ -2094,8 +2233,9 @@ static std::string rewriteScalarGMStoreFlushMarkersInFunction( size_t fallbackIndex = lines.size(); for (size_t i = lines.size(); i > 0; --i) { llvm::StringRef trimmed = llvm::StringRef(lines[i - 1]).trim(); - if (trimmed.empty()) + if (trimmed.empty()) { continue; + } if (trimmed.starts_with("}")) fallbackIndex = i - 1; break; @@ -2119,12 +2259,14 @@ static std::string rewriteScalarGMStoreFlushMarkersInFunction( inserted = true; } out.append(lines[i]); - if (i + 1 < lines.size() || hasTrailingNewline) + if (i + 1 < lines.size() || hasTrailingNewline) { out.push_back('\n'); + } } - if (!inserted) + if (!inserted) { appendScalarGMFlush(out, " "); + } return out; } @@ -2154,27 +2296,32 @@ static void rewriteScalarGMStoreFlushMarkers(std::string &cpp) { ref = split.second; llvm::StringRef trimmed = llvm::StringRef(line).trim(); - if (!inFunction && isAICOREFunctionStart(trimmed)) + if (!inFunction && isAICOREFunctionStart(trimmed)) { inFunction = true; + } if (!inFunction) { out.append(line); - if (hadNewline) + if (hadNewline) { out.push_back('\n'); + } continue; } functionLines.push_back(std::move(line)); int delta = countBraceDelta(functionLines.back()); - if (delta != 0) + if (delta != 0) { sawFunctionBrace = true; + } braceDepth += delta; - if (sawFunctionBrace && braceDepth == 0) + if (sawFunctionBrace && braceDepth == 0) { flushFunction(hadNewline); + } } - if (!functionLines.empty()) + if (!functionLines.empty()) { flushFunction(false); + } cpp.swap(out); } @@ -2196,8 +2343,9 @@ static bool isPreprocessorDirectiveLine(llvm::StringRef trimmedLine) { // Trim only those malformed suffixes here so bisheng can compile the emitted // source until the upstream printer behavior is fixed. static void rewriteMalformedVerbatimSemicolons(std::string &cpp) { - if (cpp.empty()) + if (cpp.empty()) { return; + } llvm::StringRef input(cpp); std::string rewritten; @@ -2220,8 +2368,9 @@ static void rewriteMalformedVerbatimSemicolons(std::string &cpp) { } else { if (isPreprocessorDirectiveLine(trimmed) && trimmed.ends_with(";")) { size_t semicolonPos = current.find_last_of(';'); - if (semicolonPos != std::string::npos) + if (semicolonPos != std::string::npos) { current.erase(semicolonPos, 1); + } } else if (!trimmed.empty() && !trimmed.starts_with("//") && !trimmed.starts_with("/*") && trimmed.ends_with(";;")) { size_t semicolonPos = current.find_last_of(';'); @@ -2273,10 +2422,12 @@ static bool rewriteAddPtrTraceMarkers(std::string &cpp, bool showTrace) { size_t replaceEnd = call->rparenPos; if (!showTrace) { size_t i = call->rparenPos + 1; - while (i < cpp.size() && std::isspace(static_cast(cpp[i]))) + while (i < cpp.size() && std::isspace(static_cast(cpp[i]))) { ++i; - if (i < cpp.size() && cpp[i] == ';') + } + if (i < cpp.size() && cpp[i] == ';') { replaceEnd = i; + } } cpp.replace(call->markerPos, (replaceEnd - call->markerPos) + 1, @@ -2297,11 +2448,13 @@ static bool isGeneratedGlobalTensorDecl(llvm::StringRef trimmed, decl = trimmed.drop_back().rtrim(); size_t lastWs = decl.find_last_of(" \t"); - if (lastWs == llvm::StringRef::npos) + if (lastWs == llvm::StringRef::npos) { return false; + } varName = decl.drop_front(lastWs + 1); - if (!varName.starts_with("v") || varName.size() <= 1) + if (!varName.starts_with("v") || varName.size() <= 1) { return false; + } return llvm::all_of(varName.drop_front(1), [](char c) { return std::isdigit(c); }); } @@ -2330,8 +2483,9 @@ static void rewriteHoistedGlobalTensorDecls(std::string &cpp) { llvm::StringRef varName; if (isGeneratedGlobalTensorDecl(trimmed, decl, varName)) { size_t indentLen = line.find_first_not_of(" \t"); - if (indentLen == std::string::npos) + if (indentLen == std::string::npos) { indentLen = 0; + } llvm::StringRef indent = line.take_front(indentLen); out.append(indent.str()); @@ -2340,10 +2494,12 @@ static void rewriteHoistedGlobalTensorDecls(std::string &cpp) { rewritten = true; } - if (!rewritten) + if (!rewritten) { out.append(line.str()); - if (!rest.empty()) + } + if (!rest.empty()) { out.push_back('\n'); + } ref = rest; } @@ -2354,12 +2510,15 @@ static std::optional> parseNameHintMarker(llvm::StringRef markerBody) { auto decodeHintMarkerToken = [](llvm::StringRef token) { auto hexValue = [](char c) -> int { - if (c >= '0' && c <= '9') + if (c >= '0' && c <= '9') { return c - '0'; - if (c >= 'a' && c <= 'f') + } + if (c >= 'a' && c <= 'f') { return c - 'a' + 10; - if (c >= 'A' && c <= 'F') + } + if (c >= 'A' && c <= 'F') { return c - 'A' + 10; + } return -1; }; @@ -2384,8 +2543,9 @@ parseNameHintMarker(llvm::StringRef markerBody) { llvm::SmallVector hints; markerBody = markerBody.trim(); - if (markerBody.empty()) + if (markerBody.empty()) { return std::nullopt; + } size_t start = 0; while (start <= markerBody.size()) { @@ -2393,15 +2553,18 @@ parseNameHintMarker(llvm::StringRef markerBody) { llvm::StringRef token = markerBody.slice( start, comma == llvm::StringRef::npos ? markerBody.size() : comma); token = token.trim(); - if (!token.empty()) + if (!token.empty()) { hints.push_back(decodeHintMarkerToken(token)); - if (comma == llvm::StringRef::npos) + } + if (comma == llvm::StringRef::npos) { break; + } start = comma + 1; } - if (hints.empty()) + if (hints.empty()) { return std::nullopt; + } return hints; } @@ -2497,8 +2660,9 @@ static void emitProvenanceComments(std::string &segment) { if (names && !names->empty()) { out.append("// pto: "); for (size_t idx = 0; idx < names->size(); ++idx) { - if (idx != 0) + if (idx != 0) { out.append(", "); + } out.push_back('%'); out.append(sanitizeCommentText((*names)[idx])); } @@ -2532,15 +2696,17 @@ struct ConstantDeclCandidate { } // namespace static bool isGeneratedValueName(llvm::StringRef name) { - if (!name.consume_front("v") || name.empty()) + if (!name.consume_front("v") || name.empty()) { return false; + } return llvm::all_of(name, [](char c) { return std::isdigit(c); }); } static bool isConstFoldableScalarType(llvm::StringRef type) { type = type.trim(); - if (type.starts_with("const ") || type.starts_with("constexpr ")) + if (type.starts_with("const ") || type.starts_with("constexpr ")) { return false; + } return llvm::StringSwitch(type) .Cases("bool", "float", "double", "half", "bfloat16_t", true) .Cases("int8_t", "uint8_t", "int16_t", "uint16_t", true) @@ -2550,10 +2716,12 @@ static bool isConstFoldableScalarType(llvm::StringRef type) { static bool isLiteralInitializer(llvm::StringRef rhs) { rhs = rhs.trim(); - if (rhs.empty()) + if (rhs.empty()) { return false; - if (rhs == "true" || rhs == "false" || rhs == "nullptr") + } + if (rhs == "true" || rhs == "false" || rhs == "nullptr") { return true; + } static const llvm::Regex kIntLiteral( R"(^[+-]?(0[xX][0-9A-Fa-f]+|[0-9]+)[uUlL]*$)"); @@ -2573,10 +2741,12 @@ static std::string normalizeConstInitializer(llvm::StringRef type, type = type.trim(); rhs = rhs.trim(); if (type == "bool") { - if (rhs == "0" || rhs == "false") + if (rhs == "0" || rhs == "false") { return "false"; - if (rhs == "1" || rhs == "-1" || rhs == "true") + } + if (rhs == "1" || rhs == "-1" || rhs == "true") { return "true"; + } } return rhs.str(); } @@ -2606,24 +2776,28 @@ static bool parseConstantDeclarationLine(llvm::StringRef line, } size_t lastWs = lhs.find_last_of(" \t"); - if (lastWs == llvm::StringRef::npos) + if (lastWs == llvm::StringRef::npos) { return false; + } llvm::StringRef type = lhs.take_front(lastWs).rtrim(); llvm::StringRef name = lhs.drop_front(lastWs + 1).trim(); - if (!isGeneratedValueName(name) || !isConstFoldableScalarType(type)) + if (!isGeneratedValueName(name) || !isConstFoldableScalarType(type)) { return false; + } size_t indentLen = line.find_first_not_of(" \t"); - if (indentLen == llvm::StringRef::npos) + if (indentLen == llvm::StringRef::npos) { indentLen = 0; + } candidate.indent = line.take_front(indentLen).str(); candidate.type = type.str(); valueName = name.str(); if (!rhs.empty()) { - if (!isLiteralInitializer(rhs)) + if (!isLiteralInitializer(rhs)) { return false; + } candidate.hasInitializer = true; candidate.initializer = normalizeConstInitializer(type, rhs); } @@ -2641,13 +2815,15 @@ static bool parseGeneratedValueAssignment(llvm::StringRef line, llvm::StringRef body = trimmed.drop_back().rtrim(); size_t eqPos = body.find('='); - if (eqPos == llvm::StringRef::npos) + if (eqPos == llvm::StringRef::npos) { return false; + } llvm::StringRef lhs = body.take_front(eqPos).rtrim(); rhs = body.drop_front(eqPos + 1).trim(); - if (!isGeneratedValueName(lhs)) + if (!isGeneratedValueName(lhs)) { return false; + } valueName = lhs; return true; } @@ -2674,12 +2850,14 @@ static void rewriteScalarConstantDecls(std::string &cpp) { llvm::StringRef assignedName; llvm::StringRef rhs; - if (!parseGeneratedValueAssignment(lines[i], assignedName, rhs)) + if (!parseGeneratedValueAssignment(lines[i], assignedName, rhs)) { continue; + } auto it = candidates.find(assignedName); - if (it == candidates.end()) + if (it == candidates.end()) { continue; + } ConstantDeclCandidate &info = it->second; ++info.assignmentCount; @@ -2693,14 +2871,17 @@ static void rewriteScalarConstantDecls(std::string &cpp) { std::string initializer; if (info.hasInitializer) { - if (info.assignmentCount != 0) + if (info.assignmentCount != 0) { continue; + } initializer = info.initializer; } else { - if (info.assignmentCount != 1) + if (info.assignmentCount != 1) { continue; - if (!isLiteralInitializer(info.assignmentRhs)) + } + if (!isLiteralInitializer(info.assignmentRhs)) { continue; + } initializer = normalizeConstInitializer( info.type, llvm::StringRef(info.assignmentRhs)); eraseLine[info.assignmentLine] = true; @@ -2722,20 +2903,24 @@ static void rewriteScalarConstantDecls(std::string &cpp) { --braceDepth; } - if (depthBefore == 0 && braceDepth > 0) + if (depthBefore == 0 && braceDepth > 0) { segmentStart = i; - if (depthBefore > 0 && braceDepth == 0) + } + if (depthBefore > 0 && braceDepth == 0) { rewriteSegment(segmentStart, i); + } } std::string out; out.reserve(cpp.size()); for (size_t i = 0; i < lines.size(); ++i) { - if (eraseLine[i]) + if (eraseLine[i]) { continue; + } out.append(lines[i]); - if (i + 1 != lines.size()) + if (i + 1 != lines.size()) { out.push_back('\n'); + } } cpp.swap(out); } @@ -2777,8 +2962,9 @@ static void prepareVPTOForEmission(PassManager &pm) { createVPTOExpandWrapperOpsPass()); kernelModulePM.addNestedPass( pto::createPTOInferVPTOVecScopePass()); - if (enableSoftPostUpdate) + if (enableSoftPostUpdate) { kernelModulePM.addPass(pto::createVPTOSoftPostUpdatePass()); + } kernelModulePM.addPass(createLoopInvariantCodeMotionPass()); kernelModulePM.addNestedPass( pto::createPTONarrowVPTOLoopCountersPass()); @@ -2843,8 +3029,9 @@ buildVPTOEmissionOptions(const pto::CANNVersion &cannVersion, options.targetTriple = "hiipu64-hisilicon-cce"; options.cannVersion = cannVersion; std::string arch = normalizeArch(targetArch); - if (isA2A3Arch(arch)) + if (isA2A3Arch(arch)) { options.march = "dav-c220-vec"; + } return options; } @@ -2904,8 +3091,9 @@ static LogicalResult runVPTOBackendPipeline(OwningOpRef &module, pm.enableVerifier(); pm.addPass(pto::createVPTOSplitCVModulePass()); pm.addPass(pto::createVPTONormalizeContainerPass()); - if (hasTileOpsToExpand) + if (hasTileOpsToExpand) { lowerPTOToVPTOBackend(pm, module.get()); + } auto &kernelModulePM = pm.nest(); // Inline legal direct calls before VMI layout assignment so private helper // bodies participate in one caller-local layout decision. The Func @@ -2970,8 +3158,9 @@ int mlir::pto::compilePTOASModule( // Validate stack-local struct provenance before every output path. In // particular, --emit-pto-ir returns before the EmitC validation pass and // VPTO does not use that pass. - if (failed(pto::validateStructProvenance(*module))) + if (failed(pto::validateStructProvenance(*module))) { return 1; + } std::string arch = resolveEffectiveTargetArch(*module, context.getArch()); @@ -3073,8 +3262,9 @@ int mlir::pto::compilePTOASModule( module->walk([&](mlir::func::FuncOp func) { auto hintAttr = func->getAttrOfType("pto.auto_sync_tail_hint"); - if (!hintAttr) + if (!hintAttr) { return; + } std::string normalizedHint; if (!parseAutoSyncTailHint(hintAttr.getValue(), normalizedHint)) { @@ -3088,8 +3278,9 @@ int mlir::pto::compilePTOASModule( func->setAttr("pto.auto_sync_tail_hint", mlir::StringAttr::get(module->getContext(), normalizedHint)); }); - if (invalidAutoSyncTailHint) + if (invalidAutoSyncTailHint) { return 1; + } bool hasTAssign = false; module->walk([&](pto::TAssignOp) { hasTAssign = true; }); @@ -3133,14 +3324,16 @@ int mlir::pto::compilePTOASModule( bool hasUserPlannedMultiAddrs = false; module->walk([&](pto::AllocMultiTileOp op) { - if (!op->hasAttr(pto::kPtoMultiBufferAddrsAttrName)) + if (!op->hasAttr(pto::kPtoMultiBufferAddrsAttrName)) { return; + } op.emitError() << "attribute '" << pto::kPtoMultiBufferAddrsAttrName << "' is reserved for pto-plan-memory"; hasUserPlannedMultiAddrs = true; }); - if (hasUserPlannedMultiAddrs) + if (hasUserPlannedMultiAddrs) { return 1; + } if (effectiveLevel == PTOBuildLevel::Level3) { // In level3 the caller owns local memory and PTOPlanMemory is skipped, so @@ -3160,8 +3353,9 @@ int mlir::pto::compilePTOASModule( missing = true; } }); - if (missing) + if (missing) { return 1; + } } else { bool hasAddr = false; module->walk([&](pto::AllocTileOp op) { @@ -3178,12 +3372,14 @@ int mlir::pto::compilePTOASModule( hasAddr = true; } }); - if (hasAddr) + if (hasAddr) { return 1; + } } - if (!validateReserveBufferLevelRules(*module, effectiveLevel)) + if (!validateReserveBufferLevelRules(*module, effectiveLevel)) { return 1; + } { PassManager preBackendPM(module->getContext()); @@ -3206,8 +3402,9 @@ int mlir::pto::compilePTOASModule( "skipping the shared PTO-to-VPTO lowering pipeline.\n"; return 1; } - if (failed(runVPTOBackendPipeline(module, hasTileOpsToExpand))) + if (failed(runVPTOBackendPipeline(module, hasTileOpsToExpand))) { return 1; + } return emitVPTOBackendResult(*module, result, emitVPTOHostStub, context.getCANNVersionOrDefault()); } @@ -3215,8 +3412,9 @@ int mlir::pto::compilePTOASModule( // Main PassManager PassManager pm(module->getContext()); - if (failed(applyPassManagerCLOptions(pm))) + if (failed(applyPassManagerCLOptions(pm))) { return 1; + } // Rank-2 → rank-5 view canonicalization is currently gated on the VPTO // backend to limit blast radius. A3/A5 EmitC codegen already pads strides @@ -3224,28 +3422,32 @@ int mlir::pto::compilePTOASModule( // does not need the canonicalization pass at the IR level. When VPTO // validation is complete and the pass is proven stable, the gate can be // lifted to make it unconditional for all backends. - if (effectiveBackend == PTOBackend::VPTO) + if (effectiveBackend == PTOBackend::VPTO) { pm.addNestedPass(pto::createPTOCanonicalizeIRPass()); + } pm.addPass(createSerialFrontendPipeLoweringPass()); //pm.addNestedPass(pto::createPTOVerifyTFreePass()); pm.addPass(pto::createPTOInferValidatePipeInitPass()); pm.addNestedPass(pto::createLoweringSyncToPipePass()); - if (!disableInferLayout) + if (!disableInferLayout) { pm.addNestedPass(pto::createInferPTOLayoutPass()); + } // PTOViewToMemref is generic view lowering required by both backends; keep it // outside the local-memory planning gate so default A2/A3 EmitC still lowers // pto.make_tensor_view before backend legalization. const bool isA2A3 = isA2A3Arch(arch); - if (!isA2A3) + if (!isA2A3) { pm.addNestedPass(pto::createPTOA5NormalizeTMovPass()); + } pm.addNestedPass( pto::createPTOValidateIntToPtrUsesPass()); // PTODSL legality discovery happens on tile-native PTO IR before fusion. // Fusion may later filter the ordered `candidates` array; ExpandTileOp // consumes the first candidate that remains. - if (!isA2A3 && effectiveBackend == PTOBackend::VPTO && hasTileOpsToExpand) + if (!isA2A3 && effectiveBackend == PTOBackend::VPTO && hasTileOpsToExpand) { pm.addPass(pto::createInsertTemplateAttributesPass()); + } // Keep frontend fusion on tile-native PTO IR and annotate last_use directly // on scheduled block-local spans before the shared mainline lowers tiles. @@ -3347,8 +3549,9 @@ int mlir::pto::compilePTOASModule( // Materialize each `pto.multi_tile_get` as an addressed `pto.alloc_tile`; // dynamic selections use an `arith.select` chain over planned addresses. pm.addPass(pto::createPTOResolveBufferSelectPass()); - if (effectiveBackend == PTOBackend::EmitC) + if (effectiveBackend == PTOBackend::EmitC) { pm.addPass(createNarrowUnusedMultiResultProvenancePass()); + } module->getOperation()->setAttr( "pto.target_arch", @@ -3369,12 +3572,14 @@ int mlir::pto::compilePTOASModule( pm.addPass(createCSEPass()); // PTODSL backend helpers already use the tile-native ABI. pm.addPass(pto::createPTOInlineBackendHelpersPass()); - if (effectiveBackend == PTOBackend::EmitC) + if (effectiveBackend == PTOBackend::EmitC) { pm.addPass(createNarrowUnusedMultiResultProvenancePass()); + } pm.addPass(createCanonicalizerPass()); pm.addPass(createCSEPass()); - if (failed(applyConfiguredPassManagerCLOptions(pm, "main PTOAS pipeline"))) + if (failed(applyConfiguredPassManagerCLOptions(pm, "main PTOAS pipeline"))) { return 1; + } if (effectiveBackend == PTOBackend::VPTO) { if (failed(pm.run(*module))) { @@ -3382,17 +3587,20 @@ int mlir::pto::compilePTOASModule( return 1; } - if (ptoPrintSeamIR) + if (ptoPrintSeamIR) { printSharedPreBackendSeamIR(*module); + } if (ptoPrintSeamIR) { module->print(llvm::errs()); llvm::errs() << "\n"; } - if (failed(emitSharedPreBackendSeamIR(*module, ptoSeamIRFile))) + if (failed(emitSharedPreBackendSeamIR(*module, ptoSeamIRFile))) { return 1; + } - if (failed(runVPTOBackendPipeline(module, hasTileOpsToExpand))) + if (failed(runVPTOBackendPipeline(module, hasTileOpsToExpand))) { return 1; + } return emitVPTOBackendResult(*module, result, emitVPTOHostStub, context.getCANNVersionOrDefault()); } @@ -3402,10 +3610,12 @@ int mlir::pto::compilePTOASModule( return 1; } - if (ptoPrintSeamIR) + if (ptoPrintSeamIR) { printSharedPreBackendSeamIR(*module); - if (failed(emitSharedPreBackendSeamIR(*module, ptoSeamIRFile))) + } + if (failed(emitSharedPreBackendSeamIR(*module, ptoSeamIRFile))) { return 1; + } narrowUnusedMultiResultProvenanceLocs(module.get()); splitDerivedSingleResultProvenanceLocs(module.get()); diff --git a/tools/ptobc/src/canonical_printer.cpp b/tools/ptobc/src/canonical_printer.cpp index 52c1b0815a..88020aafde 100644 --- a/tools/ptobc/src/canonical_printer.cpp +++ b/tools/ptobc/src/canonical_printer.cpp @@ -62,7 +62,9 @@ static std::string joinLines(const std::vector &lines) { std::string out; for (size_t i = 0; i < lines.size(); ++i) { out += lines[i]; - if (i + 1 < lines.size()) out.push_back('\n'); + if (i + 1 < lines.size()) { + out.push_back('\n'); + } } return out; } @@ -93,7 +95,9 @@ static std::string hexFloatLiteral(mlir::FloatAttr a) { static void sortAttributesLexicographically(mlir::ModuleOp module) { module.walk([&](mlir::Operation *op) { auto attrs = op->getAttrs(); - if (attrs.size() <= 1) return; + if (attrs.size() <= 1) { + return; + } NamedAttributeVector sorted(attrs.begin(), attrs.end()); llvm::sort(sorted, [](const mlir::NamedAttribute &a, const mlir::NamedAttribute &b) { @@ -157,15 +161,23 @@ static void collectSSADefsFromSignature(const std::string &line, // Collect `%name:` occurrences within (...) in `func.func` or `^bb` lines. // This is a lightweight heuristic but works for standard MLIR assembly. size_t lpar = line.find('('); - if (lpar == std::string::npos) return; + if (lpar == std::string::npos) { + return; + } size_t rpar = line.find(')', lpar + 1); - if (rpar == std::string::npos) return; + if (rpar == std::string::npos) { + return; + } for (size_t i = lpar + 1; i < rpar; ++i) { - if (line[i] != '%') continue; + if (line[i] != '%') { + continue; + } size_t j = i + 1; while (j < rpar && isSSAIdentChar(line[j])) ++j; - if (j == i + 1) continue; + if (j == i + 1) { + continue; + } // Must be followed by ':' to be an arg/blkarg. if (j < rpar && line[j] == ':') { out.push_back(line.substr(i + 1, j - (i + 1))); @@ -204,18 +216,22 @@ static bool findConstantDefinition(const std::vector &lines, const std::string &name, std::string &imm, std::string &ty) { for (const auto &line : lines) { - if (line.find('%' + name) == std::string::npos) + if (line.find('%' + name) == std::string::npos) { continue; - if (line.find("= arith.constant") == std::string::npos) + } + if (line.find("= arith.constant") == std::string::npos) { continue; + } size_t pos = line.find('%'); - if (pos == std::string::npos) + if (pos == std::string::npos) { continue; + } size_t end = pos + 1; while (end < line.size() && isSSAIdentChar(line[end])) ++end; - if (line.substr(pos + 1, end - (pos + 1)) != name) + if (line.substr(pos + 1, end - (pos + 1)) != name) { continue; + } return parseConstantLine(line, imm, ty); } return false; diff --git a/tools/ptobc/src/leb128.cpp b/tools/ptobc/src/leb128.cpp index 71a40259d9..45a9c6d3c5 100644 --- a/tools/ptobc/src/leb128.cpp +++ b/tools/ptobc/src/leb128.cpp @@ -19,9 +19,9 @@ namespace ptobc { namespace { constexpr unsigned kLeb128PayloadBits = 7; -constexpr uint8_t kLeb128PayloadMask = 0x7fu; -constexpr uint8_t kLeb128ContinuationBit = 0x80u; -constexpr uint8_t kLeb128SignBit = 0x40u; +constexpr uint8_t kLeb128PayloadMask = 0x7FU; +constexpr uint8_t kLeb128ContinuationBit = 0x80U; +constexpr uint8_t kLeb128SignBit = 0x40U; constexpr unsigned kInt64MaxShift = 63; constexpr unsigned kInt64BitWidth = 64; } // namespace diff --git a/tools/ptobc/src/mlir_encode.cpp b/tools/ptobc/src/mlir_encode.cpp index 7d72c84851..1d120b3a79 100644 --- a/tools/ptobc/src/mlir_encode.cpp +++ b/tools/ptobc/src/mlir_encode.cpp @@ -296,7 +296,7 @@ static void appendAPIntBytesLE(Buffer &buffer, const llvm::APInt &bits) { for (unsigned i = 0; i < byteLen; ++i) { unsigned word = i / 8; unsigned off = (i % 8) * 8; - uint8_t byte = uint8_t((words[word] >> off) & 0xFFu); + uint8_t byte = uint8_t((words[word] >> off) & 0xFFU); buffer.bytes.push_back(byte); } } diff --git a/tools/ptobc/src/ptobc_format.cpp b/tools/ptobc/src/ptobc_format.cpp index 115bcac641..75b61731a1 100644 --- a/tools/ptobc/src/ptobc_format.cpp +++ b/tools/ptobc/src/ptobc_format.cpp @@ -23,7 +23,7 @@ namespace ptobc { namespace { constexpr unsigned kBitsPerByte = 8; -constexpr uint32_t kByteMask = 0xffu; +constexpr uint32_t kByteMask = 0xffU; constexpr unsigned kU32SecondByteShift = kBitsPerByte; constexpr unsigned kU32ThirdByteShift = 2 * kBitsPerByte; constexpr unsigned kU32FourthByteShift = 3 * kBitsPerByte; From 1b5c81eaa83a809749584026ab454efe188366c6 Mon Sep 17 00:00:00 2001 From: Zhang Zhendong Date: Wed, 12 Aug 2026 12:25:50 +0800 Subject: [PATCH 097/122] Merge pull request #1137 from Zhendong404/codex/issue-1126 fix(ptodsl): normalize integer boolean branch merges (#1126) --- ptodsl/docs/user_guide/05-control-flow.md | 26 ++++ ptodsl/ptodsl/_ast_rewrite.py | 148 ++++++++++++++++++--- ptodsl/ptodsl/_control_flow.py | 40 +++++- ptodsl/ptodsl/_scalar_adaptation.py | 23 +++- ptodsl/tests/test_issue_1126_bool_merge.py | 95 +++++++++++++ ptodsl/tests/test_jit_compile.py | 85 ++++++++++++ 6 files changed, 395 insertions(+), 22 deletions(-) create mode 100644 ptodsl/tests/test_issue_1126_bool_merge.py diff --git a/ptodsl/docs/user_guide/05-control-flow.md b/ptodsl/docs/user_guide/05-control-flow.md index 71440ee721..6768a3f909 100644 --- a/ptodsl/docs/user_guide/05-control-flow.md +++ b/ptodsl/docs/user_guide/05-control-flow.md @@ -213,6 +213,11 @@ conditional closes, `br.val` is the SSA-merged result seen by downstream code. This surface avoids explicit result-type declarations and explicit `pto.yield_(...)` in user code while still keeping the merge contract explicit. +When one branch yields `pto.i1` and the other yields an integer-like scalar, +the integer value is normalized to `i1` with nonzero truthiness (`0` is false; +any nonzero value is true). Other incompatible branch result types still +require an explicit conversion. + ## 5.4 `pto.const_expr` and tracing `pto.const_expr` parameters (Section 3.6) are compile-time constants. They are fixed at `.compile()` time and cannot change between launches of the same compiled kernel. Because their values are known during tracing, they interact naturally with Python control flow: @@ -277,6 +282,27 @@ def ast_rewrite_branch_kernel(): The assigned value `total` is live after the branch, so PTODSL rewrites the branch into a `pto.if_` with automatic merge. +Static list slots can also be live across a rewritten branch. The subscript +index must be an integer that can be resolved during AST rewriting, including +compile-time constants and `pto.const_expr` values: + +```python +@pto.jit(target="a5") +def ast_rewrite_static_slot_kernel(*, SLOT: pto.const_expr = 0): + values = [pto.const(0, dtype=pto.i32)] + + if pto.const(1, dtype=pto.i1): + values[SLOT] = pto.const(1, dtype=pto.i32) + + if values[SLOT]: + pto.pipe_barrier(pto.Pipe.ALL) +``` + +The rewritten slot is merged as an `scf.if` result, just like a named scalar +value. Dynamic indices and container aliases remain unsupported; use an +explicit `pto.if_`/`pto.for_` state value or a real buffer load/store for those +cases. + If a live-out value is assigned in only one branch, PTODSL keeps the old value on the missing branch: diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index ee416cbf60..10720b0d78 100644 --- a/ptodsl/ptodsl/_ast_rewrite.py +++ b/ptodsl/ptodsl/_ast_rewrite.py @@ -972,6 +972,42 @@ def visit_Subscript(self, node): return self.generic_visit(node) +class _SlotValueRewriter(ast.NodeTransformer): + """Replace selected static list slots with scalar branch state names.""" + + def __init__(self, slot_values, static_env, static_iters=None): + self._slot_values = dict(slot_values) + self._static_env = static_env + self._static_iters = dict(static_iters or {}) + + def visit_For(self, node): + if _is_pto_attr_call(node.iter, "static_range") and isinstance(node.target, ast.Name): + values = _try_eval_static_range(node.iter, self._static_env, self._static_iters) + old = self._static_iters.get(node.target.id) + if values is not None: + self._static_iters[node.target.id] = values + try: + node.body = [self.visit(stmt) for stmt in node.body] + finally: + if values is not None: + if old is None: + self._static_iters.pop(node.target.id, None) + else: + self._static_iters[node.target.id] = old + node.orelse = [self.visit(stmt) for stmt in node.orelse] + return node + return self.generic_visit(node) + + def visit_Subscript(self, node): + slots = _resolve_subscript_slots(node, self._static_env, self._static_iters, require_static=False) + if len(slots) == 1: + slot = next(iter(slots)) + value_name = self._slot_values.get(slot) + if value_name is not None: + return ast.copy_location(_name(value_name, node.ctx), node) + return self.generic_visit(node) + + class _ControlFlowRewriter: def __init__(self, static_env=None, *, section_entry_bindings=None, section_uninitialized_aliases=None): self._static_env = dict(static_env or {}) @@ -1118,16 +1154,13 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con cond_name = self._fresh("cond") then_info = _name_info(stmt.body) else_info = _name_info(stmt.orelse) + then_slot_info = _slot_info(stmt.body, self._static_env, static_iters) + else_slot_info = _slot_info(stmt.orelse, self._static_env, static_iters) assigned_slots = ( - _slot_info(stmt.body, self._static_env, static_iters).stores - | _slot_info(stmt.orelse, self._static_env, static_iters).stores + then_slot_info.stores + | else_slot_info.stores ) - if live_after_slots & assigned_slots: - slots = ", ".join(slot.display for slot in sorted(live_after_slots & assigned_slots)) - raise PTODSLAstRewriteError( - "ast_rewrite=True does not support automatic branch merges for static subscript slots yet; " - f"rewrite {slots} with explicit scalar temporaries" - ) + merge_slots = tuple(sorted(live_after_slots & assigned_slots)) assigned_any = then_info.stores | else_info.stores merge_names = tuple(sorted(live_after & assigned_any)) old_value_names = { @@ -1137,17 +1170,18 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con } branch_live_after = set(live_after) | set(merge_names) + branch_live_after_slots = set(live_after_slots) | set(merge_slots) then_body = self.rewrite_block( stmt.body, live_after=branch_live_after, - live_after_slots=live_after_slots, + live_after_slots=branch_live_after_slots, allow_loop_control=False, static_iters=static_iters, ) else_body = self.rewrite_block( stmt.orelse, live_after=branch_live_after, - live_after_slots=live_after_slots, + live_after_slots=branch_live_after_slots, allow_loop_control=False, static_iters=static_iters, ) @@ -1158,15 +1192,45 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con ) branch_name = self._fresh("br") + slot_value_names = { + # BranchHandle deliberately rejects private attribute names. Keep + # the generated branch field public while retaining a unique + # compiler-generated local name for the rewritten slot value. + slot: ( + f"pto_ast_slot_{slot.base}_" + f"{'neg' if slot.index < 0 else ''}{abs(slot.index)}_{self._counter}" + ) + for slot in merge_slots + } + self._counter += len(slot_value_names) + old_slot_value_names = { + slot: self._fresh( + f"old_slot_{slot.base}_" + f"{'neg' if slot.index < 0 else ''}{abs(slot.index)}" + ) + for slot in merge_slots + } dynamic_then_body = copy.deepcopy(then_body) dynamic_else_body = copy.deepcopy(else_body) - if merge_names: + if slot_value_names: + dynamic_then_body = [ + _SlotValueRewriter(slot_value_names, self._static_env, static_iters).visit(stmt) + for stmt in dynamic_then_body + ] + dynamic_else_body = [ + _SlotValueRewriter(slot_value_names, self._static_env, static_iters).visit(stmt) + for stmt in dynamic_else_body + ] + if merge_names or slot_value_names: dynamic_then_body.append( self._branch_assign( branch_name, merge_names, old_value_names=old_value_names, assigned_names=then_info.stores, + slot_value_names=slot_value_names, + old_slot_value_names=old_slot_value_names, + assigned_slots=then_slot_info.stores, ) ) dynamic_else_body.append( @@ -1175,6 +1239,9 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con merge_names, old_value_names=old_value_names, assigned_names=else_info.stores, + slot_value_names=slot_value_names, + old_slot_value_names=old_slot_value_names, + assigned_slots=else_slot_info.stores, ) ) @@ -1236,6 +1303,17 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con ) for name in merge_names ) + dynamic_body.extend( + ast.Assign( + targets=[_slot_subscript(slot, ast.Store())], + value=ast.Attribute( + value=_name(branch_name), + attr=slot_value_names[slot], + ctx=ast.Load(), + ), + ) + for slot in merge_slots + ) result = [ ast.Assign( @@ -1250,6 +1328,20 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con ) for name, old_name in old_value_names.items() ) + result.extend( + ast.Assign( + targets=[_name(value_name, ast.Store())], + value=_slot_subscript(slot), + ) + for slot, value_name in slot_value_names.items() + ) + result.extend( + ast.Assign( + targets=[_name(old_value_name, ast.Store())], + value=_name(slot_value_names[slot]), + ) + for slot, old_value_name in old_slot_value_names.items() + ) result.append( ast.copy_location( ast.If( @@ -1266,18 +1358,36 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con ) return result - def _branch_assign(self, branch_name, names, *, old_value_names, assigned_names): + def _branch_assign( + self, + branch_name, + names, + *, + old_value_names, + assigned_names, + slot_value_names=(), + old_slot_value_names=(), + assigned_slots=(), + ): + keywords = [ + ast.keyword( + arg=name, + value=_name(name if name in assigned_names else old_value_names[name]), + ) + for name in names + ] + keywords.extend( + ast.keyword( + arg=value_name, + value=_name(value_name if slot in assigned_slots else old_slot_value_names[slot]), + ) + for slot, value_name in slot_value_names.items() + ) return ast.Expr( value=ast.Call( func=ast.Attribute(value=_name(branch_name), attr="assign", ctx=ast.Load()), args=[], - keywords=[ - ast.keyword( - arg=name, - value=_name(name if name in assigned_names else old_value_names[name]), - ) - for name in names - ], + keywords=keywords, ) ) diff --git a/ptodsl/ptodsl/_control_flow.py b/ptodsl/ptodsl/_control_flow.py index dd42c8a23e..299f16e5bb 100644 --- a/ptodsl/ptodsl/_control_flow.py +++ b/ptodsl/ptodsl/_control_flow.py @@ -24,6 +24,7 @@ from ._diagnostics import explicit_mode_required_with_context_error from ._runtime_index_ops import coerce_runtime_index +from ._scalar_adaptation import coerce_runtime_integer_to_i1 from ._scalar_coercion import coerce_scalar_to_type from ._surface_types import const_expr from ._tracing.active import current_session, require_active_session @@ -31,7 +32,7 @@ from ._types import _StructDescriptor from ptoas.mlir.dialects import pto as _pto, scf -from ptoas.mlir.ir import InsertionPoint +from ptoas.mlir.ir import IndexType, InsertionPoint, IntegerType # ── vecscope ────────────────────────────────────────────────────────────────── @@ -423,6 +424,11 @@ def __init__(self, cond): def __enter__(self): self._cond_value = unwrap_surface_value(self._cond) + if not _is_i1_type(self._cond_value.type) and _is_integer_like_type(self._cond_value.type): + self._cond_value = coerce_runtime_integer_to_i1( + self._cond_value, + context="pto.if_(...) condition", + ) self._tmp_if = scf.IfOp(self._cond_value, hasElse=True) self._parent_block = _find_parent_block(self._tmp_if) self._handle = BranchHandle(self) @@ -555,6 +561,8 @@ def _validate_merge_spec(self): name, then_value, else_value, + then_block=self._tmp_if.then_block, + else_block=self._tmp_if.else_block, ) if then_value.type != else_value.type: raise RuntimeError( @@ -649,11 +657,39 @@ def _is_branch_assign_literal(value) -> bool: return isinstance(value, (int, float)) and not isinstance(value, bool) -def _reconcile_branch_assignment_values(name, then_value, else_value): +def _is_i1_type(type_obj) -> bool: + return IntegerType.isinstance(type_obj) and IntegerType(type_obj).width == 1 + + +def _is_integer_like_type(type_obj) -> bool: + return IndexType.isinstance(type_obj) or IntegerType.isinstance(type_obj) + + +def _coerce_integer_to_i1_at(value, *, block, context): + with InsertionPoint(block): + return coerce_runtime_integer_to_i1(value, context=context) + + +def _reconcile_branch_assignment_values(name, then_value, else_value, *, then_block=None, else_block=None): then_is_typed = hasattr(then_value, "type") else_is_typed = hasattr(else_value, "type") if then_is_typed and else_is_typed: + then_is_i1 = _is_i1_type(then_value.type) + else_is_i1 = _is_i1_type(else_value.type) + if then_is_i1 != else_is_i1: + if then_is_i1 and _is_integer_like_type(else_value.type): + else_value = _coerce_integer_to_i1_at( + else_value, + block=else_block, + context=f"br.assign(...) else branch value for '{name}'", + ) + elif else_is_i1 and _is_integer_like_type(then_value.type): + then_value = _coerce_integer_to_i1_at( + then_value, + block=then_block, + context=f"br.assign(...) then branch value for '{name}'", + ) return then_value, else_value if then_is_typed: return then_value, coerce_scalar_to_type( diff --git a/ptodsl/ptodsl/_scalar_adaptation.py b/ptodsl/ptodsl/_scalar_adaptation.py index 5178c2d964..5796a9338a 100644 --- a/ptodsl/ptodsl/_scalar_adaptation.py +++ b/ptodsl/ptodsl/_scalar_adaptation.py @@ -121,6 +121,26 @@ def coerce_runtime_integer_value(value, target_type, *, context: str): return coerce_integer_like(value, target_type) +def coerce_runtime_integer_to_i1(value, *, context: str): + """Convert one runtime integer-like value to ``i1`` using nonzero truthiness.""" + if not hasattr(value, "type"): + raise TypeError(f"{context} expects an integer-like runtime scalar, got {value!r}") + + if IndexType.isinstance(value.type): + zero = arith.ConstantOp(IndexType.get(), 0).result + return arith.CmpIOp(arith.CmpIPredicate.ne, value, zero).result + + if not IntegerType.isinstance(value.type): + raise TypeError(f"{context} expects an integer-like runtime scalar, got {value.type}") + + signless_type = _signless_integer_type(value.type) + signless_value = _strip_integer_signedness(value) + if IntegerType(signless_type).width == 1: + return signless_value + zero = arith.ConstantOp(signless_type, 0).result + return arith.CmpIOp(arith.CmpIPredicate.ne, signless_value, zero).result + + def coerce_runtime_i1_value(value, *, context: str): """Normalize one authored bool/integer-like value/literal to signless i1.""" i1_type = IntegerType.get_signless(1) @@ -136,7 +156,7 @@ def coerce_runtime_i1_value(value, *, context: str): kind = classify_runtime_scalar_type(value.type) if kind == "float": raise TypeError(f"{context} expects a bool or integer-like scalar, got {value.type}") - return coerce_integer_like(value, i1_type) + return coerce_runtime_integer_to_i1(value, context=context) def normalize_runtime_binary_operands(lhs, rhs): @@ -272,6 +292,7 @@ def _float_bytewidth(type_obj): "classify_runtime_scalar_type", "coerce_integer_like", "coerce_runtime_i1_value", + "coerce_runtime_integer_to_i1", "coerce_runtime_index_value", "coerce_runtime_integer_value", "coerce_scalar_value_to_type", diff --git a/ptodsl/tests/test_issue_1126_bool_merge.py b/ptodsl/tests/test_issue_1126_bool_merge.py new file mode 100644 index 0000000000..2c994005d1 --- /dev/null +++ b/ptodsl/tests/test_issue_1126_bool_merge.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +from ptodsl import pto, scalar + + +@pto.jit(name="issue_1126_bool_merge", target="a5") +def issue_1126_bool_merge(flag: pto.ptr(pto.i8, "gm")) -> pto.i32: + cond = pto.const(0, dtype=pto.i1) + + if pto.const(1, dtype=pto.i32): + cond = scalar.load(flag, 0) + + if cond: + return pto.const(42, dtype=pto.i32) + return pto.const(0, dtype=pto.i32) + + +def _explicit_integer_bool_merge(value): + cond = pto.const(0, dtype=pto.i1) + with pto.if_(pto.const(1, dtype=pto.i1)) as br: + with br.then_: + br.assign(value=value) + with br.else_: + br.assign(value=cond) + return br.value + + +@pto.jit(name="issue_1126_integer_widths", target="a5") +def issue_1126_integer_widths( + i8_value: pto.i8, + i16_value: pto.i16, + i32_value: pto.i32, + i64_value: pto.i64, + ui8_value: pto.ui8, + ui16_value: pto.ui16, + ui32_value: pto.ui32, + ui64_value: pto.ui64, +) -> pto.i32: + flag_i8 = _explicit_integer_bool_merge(i8_value) + flag_i16 = _explicit_integer_bool_merge(i16_value) + flag_i32 = _explicit_integer_bool_merge(i32_value) + flag_i64 = _explicit_integer_bool_merge(i64_value) + flag_ui8 = _explicit_integer_bool_merge(ui8_value) + flag_ui16 = _explicit_integer_bool_merge(ui16_value) + flag_ui32 = _explicit_integer_bool_merge(ui32_value) + flag_ui64 = _explicit_integer_bool_merge(ui64_value) + result = pto.const(0, dtype=pto.i32) + result = scalar.select(flag_i8, pto.const(1, dtype=pto.i32), result) + result = scalar.select(flag_i16, pto.const(1, dtype=pto.i32), result) + result = scalar.select(flag_i32, pto.const(1, dtype=pto.i32), result) + result = scalar.select(flag_i64, pto.const(1, dtype=pto.i32), result) + result = scalar.select(flag_ui8, pto.const(1, dtype=pto.i32), result) + result = scalar.select(flag_ui16, pto.const(1, dtype=pto.i32), result) + result = scalar.select(flag_ui32, pto.const(1, dtype=pto.i32), result) + result = scalar.select(flag_ui64, pto.const(1, dtype=pto.i32), result) + return result + + +@pto.jit(name="issue_1126_incompatible_merge", target="a5") +def issue_1126_incompatible_merge(value: pto.i32) -> pto.i32: + with pto.if_(pto.const(1, dtype=pto.i1)) as br: + with br.then_: + br.assign(value=value) + with br.else_: + br.assign(value=pto.const(0.0, dtype=pto.f32)) + return pto.const(0, dtype=pto.i32) + + +def main(): + mlir = issue_1126_bool_merge.compile().mlir_text() + assert "arith.cmpi ne" in mlir + assert "arith.cmpi ne" in mlir and "i8" in mlir + + widths_mlir = issue_1126_integer_widths.compile().mlir_text() + assert widths_mlir.count("arith.cmpi ne") == 8 + for width in ("i8", "i16", "i32", "i64"): + assert width in widths_mlir + + try: + issue_1126_incompatible_merge.compile() + except RuntimeError as exc: + assert "type mismatch for 'value'" in str(exc) + else: + raise AssertionError("incompatible integer/float branch merge should still fail") + + +if __name__ == "__main__": + main() diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 3f77d444b0..0443de30d7 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1276,6 +1276,49 @@ def ast_if_old_value_merge_probe(): _ = merged +@pto.jit(target="a5") +def ast_if_static_subscript_slot_merge_probe(): + done = [None] + done[0] = pto.const(0, dtype=pto.i32) + condition = pto.const(1, dtype=pto.i1) + + if condition: + done[0] = pto.const(1, dtype=pto.i32) + + if done[0]: + pto.pipe_barrier(pto.Pipe.ALL) + + +_AST_BRANCH_SLOT_INDEX = 1 + + +@pto.jit(target="a5") +def ast_if_static_subscript_expression_merge_probe(*, SLOT: pto.const_expr = 1): + values = [ + pto.const(0, dtype=pto.i1), + pto.const(0, dtype=pto.i1), + ] + condition = pto.const(1, dtype=pto.i1) + + if condition: + values[_AST_BRANCH_SLOT_INDEX + (SLOT - SLOT)] = pto.const(1, dtype=pto.i1) + + if values[SLOT]: + pto.pipe_barrier(pto.Pipe.ALL) + + +@pto.jit(target="a5") +def ast_if_static_subscript_negative_index_merge_probe(): + values = [pto.const(0, dtype=pto.i1)] + condition = pto.const(1, dtype=pto.i1) + + if condition: + values[-1] = pto.const(1, dtype=pto.i1) + + if values[-1]: + pto.pipe_barrier(pto.Pipe.ALL) + + @pto.jit(target="a5") def ast_if_branch_local_temp_liveness_probe(): c0 = pto.const(0, dtype=pto.i1) @@ -5770,6 +5813,48 @@ def _enter_inline_simt_with_resource_attr(): "ast_rewrite=True old-value if merge should yield from both branches", ) + ast_if_static_subscript_slot_merge_text = ast_if_static_subscript_slot_merge_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + ast_if_static_subscript_slot_merge_text, + "AST-rewritten static subscript slot if-merge specialization", + ) + expect( + ast_if_static_subscript_slot_merge_text.count("scf.if") == 2, + "AST-rewritten static subscript slot branch merge should preserve both runtime conditionals", + ) + expect( + "-> (i32)" in ast_if_static_subscript_slot_merge_text, + "AST-rewritten static subscript slot merge should return the merged slot value", + ) + expect( + "pto.barrier " in ast_if_static_subscript_slot_merge_text, + "static subscript slot values should remain usable after the branch merge", + ) + + ast_if_static_subscript_expression_merge_text = ( + ast_if_static_subscript_expression_merge_probe.compile(SLOT=1).mlir_text() + ) + expect_parse_roundtrip_and_verify( + ast_if_static_subscript_expression_merge_text, + "AST-rewritten static subscript expression if-merge specialization", + ) + expect( + ast_if_static_subscript_expression_merge_text.count("scf.if") == 2, + "static-expression and constexpr subscript indices should participate in branch merging", + ) + + ast_if_static_subscript_negative_index_merge_text = ( + ast_if_static_subscript_negative_index_merge_probe.compile().mlir_text() + ) + expect_parse_roundtrip_and_verify( + ast_if_static_subscript_negative_index_merge_text, + "AST-rewritten negative static subscript if-merge specialization", + ) + expect( + ast_if_static_subscript_negative_index_merge_text.count("scf.if") == 2, + "negative static subscript indices should participate in branch merging", + ) + ast_if_branch_local_temp_liveness_text = ast_if_branch_local_temp_liveness_probe.compile().mlir_text() expect_parse_roundtrip_and_verify( ast_if_branch_local_temp_liveness_text, From 72b782624a2a223686e06262117d68df00a25445 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Tue, 11 Aug 2026 15:51:23 +0800 Subject: [PATCH 098/122] fix(vmi): preserve rounding for float-to-int vcvt (#585) --- include/PTO/IR/VMIOps.td | 9 +++-- lib/PTO/IR/VMI.cpp | 33 +++++++++++++++---- .../Transforms/VMILowerUnifiedToLegacy.cpp | 8 +++-- lib/PTO/Transforms/VMIToVPTO.cpp | 8 +++-- .../vmi_new/vmi_to_vpto_fptosi_same_width.pto | 17 ++++++++++ .../vmi_vcvt_fptosi_lower_to_legacy_new.pto | 9 +++++ 6 files changed, 70 insertions(+), 14 deletions(-) diff --git a/include/PTO/IR/VMIOps.td b/include/PTO/IR/VMIOps.td index 19399252f4..f2bc1a0364 100644 --- a/include/PTO/IR/VMIOps.td +++ b/include/PTO/IR/VMIOps.td @@ -579,6 +579,7 @@ def VMITruncFOp : VMI_Op<"truncf", [Pure]> { def VMIFPToSIOp : VMI_Op<"fptosi", [Pure]> { let summary = "VMI floating-point to signed integer elementwise conversion"; let arguments = (ins VMI_VRegTypeConstraint:$source, + OptionalAttr:$rounding, OptionalAttr:$saturate); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; @@ -588,6 +589,7 @@ def VMIFPToSIOp : VMI_Op<"fptosi", [Pure]> { def VMIFPToUIOp : VMI_Op<"fptoui", [Pure]> { let summary = "VMI floating-point to unsigned integer elementwise conversion"; let arguments = (ins VMI_VRegTypeConstraint:$source, + OptionalAttr:$rounding, OptionalAttr:$saturate); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; @@ -1381,9 +1383,10 @@ def VMICvtOp : VMI_Op<"vcvt", [Pure]> { - int → int, |dst| < |src|: integer truncation (replaces trunci) Attributes: - - `rounding`: rounding mode for fp narrowing (R=nearest-even, - A=away-from-zero, H=half-up, Z=toward-zero). Valid only when dst - bit-width < src bit-width for fp types. + - `rounding`: rounding mode for fp narrowing and fp-to-integer + conversions. Narrowing accepts R=nearest-even, A=away-from-zero, + H=half-up, Z=toward-zero; fp-to-integer accepts the hardware modes + R/A/F/C/Z. - `saturate`: "SAT" or "NOSAT"; required for fp narrowing, integer narrowing, and fp-to-int conversion. - `pmode`: predication mode ("merge" | "zero"). diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp index bc9697c2e1..8e174e5e29 100644 --- a/lib/PTO/IR/VMI.cpp +++ b/lib/PTO/IR/VMI.cpp @@ -1996,6 +1996,12 @@ LogicalResult VMIFPToSIOp::verify() { if (!contract) { return emitOpError("unsupported fp-to-si conversion element type pair"); } + if (auto roundingAttr = (*this)->getAttrOfType("rounding")) { + StringRef rounding = roundingAttr.getValue(); + if (rounding != "R" && rounding != "A" && rounding != "F" && + rounding != "C" && rounding != "Z") + return emitOpError("rounding attr must be R, A, F, C, or Z"); + } if (contract->requiresSat) { auto satAttr = (*this)->getAttrOfType("saturate"); if (!satAttr) @@ -2031,6 +2037,12 @@ LogicalResult VMIFPToUIOp::verify() { if (!contract) { return emitOpError("unsupported fp-to-ui conversion element type pair"); } + if (auto roundingAttr = (*this)->getAttrOfType("rounding")) { + StringRef rounding = roundingAttr.getValue(); + if (rounding != "R" && rounding != "A" && rounding != "F" && + rounding != "C" && rounding != "Z") + return emitOpError("rounding attr must be R, A, F, C, or Z"); + } if (contract->requiresSat) { auto satAttr = (*this)->getAttrOfType("saturate"); if (!satAttr) @@ -3938,14 +3950,21 @@ LogicalResult VMICvtOp::verify() { // --- rounding --- if (auto roundingAttr = (*this)->getAttrOfType("rounding")) { - if (dir != CvtDirection::FpNarrow) - return emitOpError("'rounding' attribute is only valid for " - "fp-narrowing conversions"); + if (dir != CvtDirection::FpNarrow && dir != CvtDirection::FpToSi && + dir != CvtDirection::FpToUi) + return emitOpError("'rounding' attribute is only valid for floating-point " + "narrowing or floating-point-to-integer conversions"); StringRef rnd = roundingAttr.getValue(); - if (rnd != "R" && rnd != "A" && rnd != "H" && rnd != "Z") - return emitOpError("rounding must be 'R' (nearest-even), " - "'A' (away-from-zero), 'H' (half-up), " - "or 'Z' (toward-zero)"); + if (dir == CvtDirection::FpNarrow) { + if (rnd != "R" && rnd != "A" && rnd != "H" && rnd != "Z") + return emitOpError("rounding must be 'R' (nearest-even), " + "'A' (away-from-zero), 'H' (half-up), " + "or 'Z' (toward-zero)"); + } else if (rnd != "R" && rnd != "A" && rnd != "F" && rnd != "C" && + rnd != "Z") { + return emitOpError("rounding must be 'R', 'A', 'F', 'C', or 'Z' for " + "floating-point-to-integer conversions"); + } } // --- saturate --- diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp index e038e01238..4a305c90fe 100644 --- a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp +++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp @@ -396,11 +396,15 @@ static LogicalResult lowerVCvt(VMICvtOp op, OpBuilder &builder) { .getResult(); } else if (direction == "fptosi") { result = - builder.create(loc, resultType, source, saturateAttr) + builder + .create(loc, resultType, source, + op.getRoundingAttr(), saturateAttr) .getResult(); } else if (direction == "fptoui") { result = - builder.create(loc, resultType, source, saturateAttr) + builder + .create(loc, resultType, source, + op.getRoundingAttr(), saturateAttr) .getResult(); } else if (direction == "sitofp") { result = diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 52253bd2ec..c43cd09410 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -12017,7 +12017,9 @@ struct OneToNVMIFPToSIOpPattern : OpConversionPattern { resultVRegTypes.push_back(resultType); } - StringAttr rnd = rewriter.getStringAttr("R"); + StringAttr rnd = op->getAttrOfType("rounding"); + if (!rnd) + rnd = rewriter.getStringAttr("R"); StringAttr sat = contract->requiresSat ? op->getAttrOfType("saturate") @@ -12230,7 +12232,9 @@ struct OneToNVMIFPToUIOpPattern : OpConversionPattern { resultVRegTypes.push_back(resultType); } - StringAttr rnd = rewriter.getStringAttr("R"); + StringAttr rnd = op->getAttrOfType("rounding"); + if (!rnd) + rnd = rewriter.getStringAttr("R"); StringAttr sat = contract->requiresSat ? op->getAttrOfType("saturate") : nullptr; diff --git a/test/lit/vmi_new/vmi_to_vpto_fptosi_same_width.pto b/test/lit/vmi_new/vmi_to_vpto_fptosi_same_width.pto index 366a571d0a..016d4a69fc 100644 --- a/test/lit/vmi_new/vmi_to_vpto_fptosi_same_width.pto +++ b/test/lit/vmi_new/vmi_to_vpto_fptosi_same_width.pto @@ -36,6 +36,19 @@ module { -> !pto.vreg<64xsi32> return %r : !pto.vreg<64xsi32> } + + // ROUND_Z must survive both unified-to-legacy and VMI-to-VPTO lowering. + func.func @f32_to_s32_round_z( + %input: !pto.vmi.vreg<64xf32, #pto.vmi.layout>) + -> !pto.vreg<64xsi32> { + %cvt = pto.vmi.vcvt %input {rounding = "Z", saturate = "SAT"} + : !pto.vmi.vreg<64xf32, #pto.vmi.layout> + -> !pto.vmi.vreg<64xsi32, #pto.vmi.layout> + %r = "pto.vmi.unpack"(%cvt) + : (!pto.vmi.vreg<64xsi32, #pto.vmi.layout>) + -> !pto.vreg<64xsi32> + return %r : !pto.vreg<64xsi32> + } } // CHECK-LABEL: func.func @f16_to_s16( @@ -45,3 +58,7 @@ module { // CHECK-LABEL: func.func @f32_to_s32( // CHECK: pto.vcvt {{.*}} {rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xsi32> // CHECK-NOT: part + +// CHECK-LABEL: func.func @f32_to_s32_round_z( +// CHECK: pto.vcvt {{.*}} {rnd = "Z", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xsi32> +// CHECK-NOT: part diff --git a/test/lit/vmi_new/vmi_vcvt_fptosi_lower_to_legacy_new.pto b/test/lit/vmi_new/vmi_vcvt_fptosi_lower_to_legacy_new.pto index eeb8b010ad..ee955aa6c6 100644 --- a/test/lit/vmi_new/vmi_vcvt_fptosi_lower_to_legacy_new.pto +++ b/test/lit/vmi_new/vmi_vcvt_fptosi_lower_to_legacy_new.pto @@ -22,6 +22,12 @@ module { %r = pto.vmi.vcvt %s {saturate = "SAT"} : !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xsi16> return %r : !pto.vmi.vreg<64xsi16> } + // Rounding must be preserved when vcvt lowers to the legacy fptosi op. + func.func @f32_to_s32_round_z(%s: !pto.vmi.vreg<64xf32>) -> !pto.vmi.vreg<64xsi32> { + %r = pto.vmi.vcvt %s {rounding = "Z", saturate = "SAT"} + : !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xsi32> + return %r : !pto.vmi.vreg<64xsi32> + } // f16->s8: NOSAT on vcvt -> NOSAT on fptosi func.func @f16_to_s8(%s: !pto.vmi.vreg<64xf16>) -> !pto.vmi.vreg<64xsi8> { %r = pto.vmi.vcvt %s {saturate = "NOSAT"} : !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xsi8> @@ -36,5 +42,8 @@ module { // CHECK-LABEL: func.func @f32_to_s16 // CHECK: pto.vmi.fptosi %{{.*}} {saturate = "SAT"} : !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xsi16> +// CHECK-LABEL: func.func @f32_to_s32_round_z +// CHECK: pto.vmi.fptosi %{{.*}} {rounding = "Z", saturate = "SAT"} : !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xsi32> + // CHECK-LABEL: func.func @f16_to_s8 // CHECK: pto.vmi.fptosi %{{.*}} {saturate = "NOSAT"} : !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xsi8> From 2c2747634fbe3b4ad8c6bcb66e7be698159a8655 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Tue, 11 Aug 2026 18:18:40 +0800 Subject: [PATCH 099/122] test(vmi): update vcvt rounding diagnostic --- test/lit/vmi_new/vmi_conversion_contract_matrix.pto | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/test/lit/vmi_new/vmi_conversion_contract_matrix.pto b/test/lit/vmi_new/vmi_conversion_contract_matrix.pto index 858a748860..88be81c7ab 100644 --- a/test/lit/vmi_new/vmi_conversion_contract_matrix.pto +++ b/test/lit/vmi_new/vmi_conversion_contract_matrix.pto @@ -1,3 +1,11 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + // RUN: pto-test-opt %s -verify-diagnostics -split-input-file // Positive matrix: six directions, all rounding tokens, saturation values, @@ -23,7 +31,7 @@ module { module { func.func @rounding_wrong_direction(%x: !pto.vmi.vreg<64xf16>) { - // expected-error@+1 {{'rounding' attribute is only valid for fp-narrowing conversions}} + // expected-error@+1 {{'rounding' attribute is only valid for floating-point narrowing or floating-point-to-integer conversions}} %r = pto.vmi.vcvt %x {rounding = "R"} : !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf32> return } From 1cecb6a19a2e359d8fe41c9c168560ec09209f53 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Wed, 12 Aug 2026 15:36:29 +0800 Subject: [PATCH 100/122] Revert "Merge pull request #1217 from hw-native-sys/codex/fix-issue-585-vmi-rounding-z" This reverts commit f71c72fee3e946b3bb87f238530e04bb10946a6f, reversing changes made to 1b5c81eaa83a809749584026ab454efe188366c6. --- include/PTO/IR/VMIOps.td | 9 ++--- lib/PTO/IR/VMI.cpp | 33 ++++--------------- .../Transforms/VMILowerUnifiedToLegacy.cpp | 8 ++--- lib/PTO/Transforms/VMIToVPTO.cpp | 8 ++--- .../vmi_conversion_contract_matrix.pto | 10 +----- .../vmi_new/vmi_to_vpto_fptosi_same_width.pto | 17 ---------- .../vmi_vcvt_fptosi_lower_to_legacy_new.pto | 9 ----- 7 files changed, 15 insertions(+), 79 deletions(-) diff --git a/include/PTO/IR/VMIOps.td b/include/PTO/IR/VMIOps.td index f2bc1a0364..19399252f4 100644 --- a/include/PTO/IR/VMIOps.td +++ b/include/PTO/IR/VMIOps.td @@ -579,7 +579,6 @@ def VMITruncFOp : VMI_Op<"truncf", [Pure]> { def VMIFPToSIOp : VMI_Op<"fptosi", [Pure]> { let summary = "VMI floating-point to signed integer elementwise conversion"; let arguments = (ins VMI_VRegTypeConstraint:$source, - OptionalAttr:$rounding, OptionalAttr:$saturate); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; @@ -589,7 +588,6 @@ def VMIFPToSIOp : VMI_Op<"fptosi", [Pure]> { def VMIFPToUIOp : VMI_Op<"fptoui", [Pure]> { let summary = "VMI floating-point to unsigned integer elementwise conversion"; let arguments = (ins VMI_VRegTypeConstraint:$source, - OptionalAttr:$rounding, OptionalAttr:$saturate); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; @@ -1383,10 +1381,9 @@ def VMICvtOp : VMI_Op<"vcvt", [Pure]> { - int → int, |dst| < |src|: integer truncation (replaces trunci) Attributes: - - `rounding`: rounding mode for fp narrowing and fp-to-integer - conversions. Narrowing accepts R=nearest-even, A=away-from-zero, - H=half-up, Z=toward-zero; fp-to-integer accepts the hardware modes - R/A/F/C/Z. + - `rounding`: rounding mode for fp narrowing (R=nearest-even, + A=away-from-zero, H=half-up, Z=toward-zero). Valid only when dst + bit-width < src bit-width for fp types. - `saturate`: "SAT" or "NOSAT"; required for fp narrowing, integer narrowing, and fp-to-int conversion. - `pmode`: predication mode ("merge" | "zero"). diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp index 8e174e5e29..bc9697c2e1 100644 --- a/lib/PTO/IR/VMI.cpp +++ b/lib/PTO/IR/VMI.cpp @@ -1996,12 +1996,6 @@ LogicalResult VMIFPToSIOp::verify() { if (!contract) { return emitOpError("unsupported fp-to-si conversion element type pair"); } - if (auto roundingAttr = (*this)->getAttrOfType("rounding")) { - StringRef rounding = roundingAttr.getValue(); - if (rounding != "R" && rounding != "A" && rounding != "F" && - rounding != "C" && rounding != "Z") - return emitOpError("rounding attr must be R, A, F, C, or Z"); - } if (contract->requiresSat) { auto satAttr = (*this)->getAttrOfType("saturate"); if (!satAttr) @@ -2037,12 +2031,6 @@ LogicalResult VMIFPToUIOp::verify() { if (!contract) { return emitOpError("unsupported fp-to-ui conversion element type pair"); } - if (auto roundingAttr = (*this)->getAttrOfType("rounding")) { - StringRef rounding = roundingAttr.getValue(); - if (rounding != "R" && rounding != "A" && rounding != "F" && - rounding != "C" && rounding != "Z") - return emitOpError("rounding attr must be R, A, F, C, or Z"); - } if (contract->requiresSat) { auto satAttr = (*this)->getAttrOfType("saturate"); if (!satAttr) @@ -3950,21 +3938,14 @@ LogicalResult VMICvtOp::verify() { // --- rounding --- if (auto roundingAttr = (*this)->getAttrOfType("rounding")) { - if (dir != CvtDirection::FpNarrow && dir != CvtDirection::FpToSi && - dir != CvtDirection::FpToUi) - return emitOpError("'rounding' attribute is only valid for floating-point " - "narrowing or floating-point-to-integer conversions"); + if (dir != CvtDirection::FpNarrow) + return emitOpError("'rounding' attribute is only valid for " + "fp-narrowing conversions"); StringRef rnd = roundingAttr.getValue(); - if (dir == CvtDirection::FpNarrow) { - if (rnd != "R" && rnd != "A" && rnd != "H" && rnd != "Z") - return emitOpError("rounding must be 'R' (nearest-even), " - "'A' (away-from-zero), 'H' (half-up), " - "or 'Z' (toward-zero)"); - } else if (rnd != "R" && rnd != "A" && rnd != "F" && rnd != "C" && - rnd != "Z") { - return emitOpError("rounding must be 'R', 'A', 'F', 'C', or 'Z' for " - "floating-point-to-integer conversions"); - } + if (rnd != "R" && rnd != "A" && rnd != "H" && rnd != "Z") + return emitOpError("rounding must be 'R' (nearest-even), " + "'A' (away-from-zero), 'H' (half-up), " + "or 'Z' (toward-zero)"); } // --- saturate --- diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp index 4a305c90fe..e038e01238 100644 --- a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp +++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp @@ -396,15 +396,11 @@ static LogicalResult lowerVCvt(VMICvtOp op, OpBuilder &builder) { .getResult(); } else if (direction == "fptosi") { result = - builder - .create(loc, resultType, source, - op.getRoundingAttr(), saturateAttr) + builder.create(loc, resultType, source, saturateAttr) .getResult(); } else if (direction == "fptoui") { result = - builder - .create(loc, resultType, source, - op.getRoundingAttr(), saturateAttr) + builder.create(loc, resultType, source, saturateAttr) .getResult(); } else if (direction == "sitofp") { result = diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index c43cd09410..52253bd2ec 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -12017,9 +12017,7 @@ struct OneToNVMIFPToSIOpPattern : OpConversionPattern { resultVRegTypes.push_back(resultType); } - StringAttr rnd = op->getAttrOfType("rounding"); - if (!rnd) - rnd = rewriter.getStringAttr("R"); + StringAttr rnd = rewriter.getStringAttr("R"); StringAttr sat = contract->requiresSat ? op->getAttrOfType("saturate") @@ -12232,9 +12230,7 @@ struct OneToNVMIFPToUIOpPattern : OpConversionPattern { resultVRegTypes.push_back(resultType); } - StringAttr rnd = op->getAttrOfType("rounding"); - if (!rnd) - rnd = rewriter.getStringAttr("R"); + StringAttr rnd = rewriter.getStringAttr("R"); StringAttr sat = contract->requiresSat ? op->getAttrOfType("saturate") : nullptr; diff --git a/test/lit/vmi_new/vmi_conversion_contract_matrix.pto b/test/lit/vmi_new/vmi_conversion_contract_matrix.pto index 88be81c7ab..858a748860 100644 --- a/test/lit/vmi_new/vmi_conversion_contract_matrix.pto +++ b/test/lit/vmi_new/vmi_conversion_contract_matrix.pto @@ -1,11 +1,3 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - // RUN: pto-test-opt %s -verify-diagnostics -split-input-file // Positive matrix: six directions, all rounding tokens, saturation values, @@ -31,7 +23,7 @@ module { module { func.func @rounding_wrong_direction(%x: !pto.vmi.vreg<64xf16>) { - // expected-error@+1 {{'rounding' attribute is only valid for floating-point narrowing or floating-point-to-integer conversions}} + // expected-error@+1 {{'rounding' attribute is only valid for fp-narrowing conversions}} %r = pto.vmi.vcvt %x {rounding = "R"} : !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf32> return } diff --git a/test/lit/vmi_new/vmi_to_vpto_fptosi_same_width.pto b/test/lit/vmi_new/vmi_to_vpto_fptosi_same_width.pto index 016d4a69fc..366a571d0a 100644 --- a/test/lit/vmi_new/vmi_to_vpto_fptosi_same_width.pto +++ b/test/lit/vmi_new/vmi_to_vpto_fptosi_same_width.pto @@ -36,19 +36,6 @@ module { -> !pto.vreg<64xsi32> return %r : !pto.vreg<64xsi32> } - - // ROUND_Z must survive both unified-to-legacy and VMI-to-VPTO lowering. - func.func @f32_to_s32_round_z( - %input: !pto.vmi.vreg<64xf32, #pto.vmi.layout>) - -> !pto.vreg<64xsi32> { - %cvt = pto.vmi.vcvt %input {rounding = "Z", saturate = "SAT"} - : !pto.vmi.vreg<64xf32, #pto.vmi.layout> - -> !pto.vmi.vreg<64xsi32, #pto.vmi.layout> - %r = "pto.vmi.unpack"(%cvt) - : (!pto.vmi.vreg<64xsi32, #pto.vmi.layout>) - -> !pto.vreg<64xsi32> - return %r : !pto.vreg<64xsi32> - } } // CHECK-LABEL: func.func @f16_to_s16( @@ -58,7 +45,3 @@ module { // CHECK-LABEL: func.func @f32_to_s32( // CHECK: pto.vcvt {{.*}} {rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xsi32> // CHECK-NOT: part - -// CHECK-LABEL: func.func @f32_to_s32_round_z( -// CHECK: pto.vcvt {{.*}} {rnd = "Z", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xsi32> -// CHECK-NOT: part diff --git a/test/lit/vmi_new/vmi_vcvt_fptosi_lower_to_legacy_new.pto b/test/lit/vmi_new/vmi_vcvt_fptosi_lower_to_legacy_new.pto index ee955aa6c6..eeb8b010ad 100644 --- a/test/lit/vmi_new/vmi_vcvt_fptosi_lower_to_legacy_new.pto +++ b/test/lit/vmi_new/vmi_vcvt_fptosi_lower_to_legacy_new.pto @@ -22,12 +22,6 @@ module { %r = pto.vmi.vcvt %s {saturate = "SAT"} : !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xsi16> return %r : !pto.vmi.vreg<64xsi16> } - // Rounding must be preserved when vcvt lowers to the legacy fptosi op. - func.func @f32_to_s32_round_z(%s: !pto.vmi.vreg<64xf32>) -> !pto.vmi.vreg<64xsi32> { - %r = pto.vmi.vcvt %s {rounding = "Z", saturate = "SAT"} - : !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xsi32> - return %r : !pto.vmi.vreg<64xsi32> - } // f16->s8: NOSAT on vcvt -> NOSAT on fptosi func.func @f16_to_s8(%s: !pto.vmi.vreg<64xf16>) -> !pto.vmi.vreg<64xsi8> { %r = pto.vmi.vcvt %s {saturate = "NOSAT"} : !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xsi8> @@ -42,8 +36,5 @@ module { // CHECK-LABEL: func.func @f32_to_s16 // CHECK: pto.vmi.fptosi %{{.*}} {saturate = "SAT"} : !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xsi16> -// CHECK-LABEL: func.func @f32_to_s32_round_z -// CHECK: pto.vmi.fptosi %{{.*}} {rounding = "Z", saturate = "SAT"} : !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xsi32> - // CHECK-LABEL: func.func @f16_to_s8 // CHECK: pto.vmi.fptosi %{{.*}} {saturate = "NOSAT"} : !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xsi8> From 4253ba2a9d5d0c99736f16223d44b8b346a505f4 Mon Sep 17 00:00:00 2001 From: likai00 Date: Fri, 7 Aug 2026 10:40:49 +0800 Subject: [PATCH 101/122] vmi.vcvt support f4x2 --- include/PTO/IR/PTOTypeDefs.td | 9 + include/PTO/IR/PTOTypeUtils.h | 1 + include/PTO/IR/VMIUtils.h | 19 +- include/pto-c/Dialect/PTO.h | 2 + lib/Bindings/Python/PTOModule.cpp | 11 + lib/CAPI/Dialect/PTO.cpp | 8 + lib/PTO/IR/PTO.cpp | 19 ++ lib/PTO/IR/PTOTypeUtils.cpp | 10 +- lib/PTO/IR/VMI.cpp | 264 ++++++++++++++---- lib/PTO/IR/VPTO.cpp | 12 + lib/PTO/Transforms/VMIToVPTO.cpp | 114 ++++++-- lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp | 34 +++ lib/PTO/Transforms/VPTOLLVMEmitter.cpp | 37 +++ ptodsl/ptodsl/_types.py | 17 +- ptodsl/ptodsl/_vmi_namespace.py | 101 ++++++- ptodsl/tests/test_jit_compile.py | 235 ++++++++++++++++ python/pto/dialects/pto.py | 2 + test/lit/pto/low_precision_type_roundtrip.pto | 4 +- .../vmi_new/vmi_bf16x2_compute_invalid.pto | 99 +++++++ .../vmi_bf16x2_conversion_pairs_invalid.pto | 129 +++++++++ ...i_bitcast_bf16_d4_to_bf16x2_d4_invalid.pto | 27 ++ .../vmi_bitcast_bf16x2_total_bits_invalid.pto | 20 ++ ...vmi_packed_fp_conversion_pairs_invalid.pto | 49 ++++ test/lit/vmi_new/vmi_to_vpto_bitcast.pto | 32 +++ .../vmi_to_vpto_extf_f4x2_to_bf16x2_ls4.pto | 38 +++ ...i_to_vpto_extf_f4x2_to_bf16x2_variants.pto | 70 +++++ ..._to_vpto_truncf_bf16x2_d2_dynamic_mask.pto | 71 +++++ ..._to_vpto_truncf_bf16x2_d2_lane_stride2.pto | 42 +++ ...mi_to_vpto_truncf_bf16x2_d2_multichunk.pto | 57 ++++ ..._to_vpto_truncf_bf16x2_d4_dynamic_mask.pto | 81 ++++++ ...mi_to_vpto_truncf_bf16x2_d4_multichunk.pto | 59 ++++ .../vmi_to_vpto_truncf_bf16x2_d4_packed4.pto | 38 +++ ...vmi_to_vpto_truncf_bf16x2_dynamic_tail.pto | 73 +++++ ...to_vpto_truncf_bf16x2_mask_granularity.pto | 35 +++ ..._vpto_truncf_bf16x2_to_f4x2_contiguous.pto | 44 +++ ...to_vpto_truncf_bf16x2_to_f4x2_variants.pto | 123 ++++++++ ...vcvt_bf16x2_f4x2_lane_mismatch_invalid.pto | 21 ++ .../vmi_vcvt_bf16x2_f4x2_rounding_invalid.pto | 67 +++++ .../vmi_vcvt_bf16x2_f4x2_saturate_invalid.pto | 32 +++ ...i_vcvt_f4x2_bf16x2_widen_attrs_invalid.pto | 35 +++ .../vmi_new/vmi_vcvt_f4x2_to_bf16_invalid.pto | 24 ++ .../vpto/bf16x2_vreg_and_bitcast_verify.pto | 31 ++ .../lit/vpto/vlogic_bf16x2_verify_invalid.pto | 46 +++ .../vpto/vmi_bf16x2_direct_load_vcvt_llvm.pto | 68 +++++ .../lit/vpto/vmi_f4x2_to_bf16x2_vcvt_llvm.pto | 69 +++++ ...i_fp4_e1_packed_surface_verify_invalid.pto | 19 +- .../vmi_fp4_packed_surface_verify_invalid.pto | 19 +- test/python/low_precision_types.py | 2 + .../compare.py | 34 +++ .../golden.py | 96 +++++++ .../kernel.pto | 50 ++++ .../launch.cpp | 40 +++ .../main.cpp | 79 ++++++ .../ptoas.flags | 1 + .../dequant-f4x2-to-bf16x2-tail/compare.py | 34 +++ .../dequant-f4x2-to-bf16x2-tail/golden.py | 99 +++++++ .../dequant-f4x2-to-bf16x2-tail/kernel.pto | 59 ++++ .../dequant-f4x2-to-bf16x2-tail/launch.cpp | 40 +++ .../dequant-f4x2-to-bf16x2-tail/main.cpp | 79 ++++++ .../dequant-f4x2-to-bf16x2-tail/ptoas.flags | 1 + .../quant-bf16x2-to-f4e2m1x2/compare.py | 34 +++ .../quant-bf16x2-to-f4e2m1x2/golden.py | 88 ++++++ .../quant-bf16x2-to-f4e2m1x2/kernel.pto | 59 ++++ .../quant-bf16x2-to-f4e2m1x2/launch.cpp | 40 +++ .../vmi_new/quant-bf16x2-to-f4e2m1x2/main.cpp | 79 ++++++ .../quant-bf16x2-to-f4e2m1x2/ptoas.flags | 1 + .../compare.py | 34 +++ .../quant-bf16x2-to-f4x2-contiguous/golden.py | 91 ++++++ .../kernel.pto | 57 ++++ .../launch.cpp | 40 +++ .../quant-bf16x2-to-f4x2-contiguous/main.cpp | 79 ++++++ .../ptoas.flags | 1 + .../quant-bf16x2-to-f4x2-full/compare.py | 34 +++ .../quant-bf16x2-to-f4x2-full/golden.py | 92 ++++++ .../quant-bf16x2-to-f4x2-full/kernel.pto | 58 ++++ .../quant-bf16x2-to-f4x2-full/launch.cpp | 40 +++ .../quant-bf16x2-to-f4x2-full/main.cpp | 79 ++++++ .../quant-bf16x2-to-f4x2-full/ptoas.flags | 1 + .../quant-bf16x2-to-f4x2-overflow/compare.py | 34 +++ .../quant-bf16x2-to-f4x2-overflow/golden.py | 86 ++++++ .../quant-bf16x2-to-f4x2-overflow/kernel.pto | 59 ++++ .../quant-bf16x2-to-f4x2-overflow/launch.cpp | 40 +++ .../quant-bf16x2-to-f4x2-overflow/main.cpp | 79 ++++++ .../quant-bf16x2-to-f4x2-overflow/ptoas.flags | 1 + .../quant-bf16x2-to-f4x2-rounding/compare.py | 34 +++ .../quant-bf16x2-to-f4x2-rounding/golden.py | 112 ++++++++ .../quant-bf16x2-to-f4x2-rounding/kernel.pto | 85 ++++++ .../quant-bf16x2-to-f4x2-rounding/launch.cpp | 40 +++ .../quant-bf16x2-to-f4x2-rounding/main.cpp | 79 ++++++ .../quant-bf16x2-to-f4x2-rounding/ptoas.flags | 1 + .../quant-bf16x2-to-f4x2-tail/compare.py | 34 +++ .../quant-bf16x2-to-f4x2-tail/golden.py | 90 ++++++ .../quant-bf16x2-to-f4x2-tail/kernel.pto | 70 +++++ .../quant-bf16x2-to-f4x2-tail/launch.cpp | 40 +++ .../quant-bf16x2-to-f4x2-tail/main.cpp | 79 ++++++ .../quant-bf16x2-to-f4x2-tail/ptoas.flags | 1 + 96 files changed, 4786 insertions(+), 116 deletions(-) create mode 100644 test/lit/vmi_new/vmi_bf16x2_compute_invalid.pto create mode 100644 test/lit/vmi_new/vmi_bf16x2_conversion_pairs_invalid.pto create mode 100644 test/lit/vmi_new/vmi_bitcast_bf16_d4_to_bf16x2_d4_invalid.pto create mode 100644 test/lit/vmi_new/vmi_bitcast_bf16x2_total_bits_invalid.pto create mode 100644 test/lit/vmi_new/vmi_packed_fp_conversion_pairs_invalid.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_extf_f4x2_to_bf16x2_ls4.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_extf_f4x2_to_bf16x2_variants.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d2_dynamic_mask.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d2_lane_stride2.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d2_multichunk.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d4_dynamic_mask.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d4_multichunk.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d4_packed4.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_dynamic_tail.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_mask_granularity.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_to_f4x2_contiguous.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_to_f4x2_variants.pto create mode 100644 test/lit/vmi_new/vmi_vcvt_bf16x2_f4x2_lane_mismatch_invalid.pto create mode 100644 test/lit/vmi_new/vmi_vcvt_bf16x2_f4x2_rounding_invalid.pto create mode 100644 test/lit/vmi_new/vmi_vcvt_bf16x2_f4x2_saturate_invalid.pto create mode 100644 test/lit/vmi_new/vmi_vcvt_f4x2_bf16x2_widen_attrs_invalid.pto create mode 100644 test/lit/vmi_new/vmi_vcvt_f4x2_to_bf16_invalid.pto create mode 100644 test/lit/vpto/bf16x2_vreg_and_bitcast_verify.pto create mode 100644 test/lit/vpto/vlogic_bf16x2_verify_invalid.pto create mode 100644 test/lit/vpto/vmi_bf16x2_direct_load_vcvt_llvm.pto create mode 100644 test/lit/vpto/vmi_f4x2_to_bf16x2_vcvt_llvm.pto create mode 100644 test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/compare.py create mode 100644 test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/golden.py create mode 100644 test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/kernel.pto create mode 100644 test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/launch.cpp create mode 100644 test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/main.cpp create mode 100644 test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/ptoas.flags create mode 100644 test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/compare.py create mode 100644 test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/golden.py create mode 100644 test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/kernel.pto create mode 100644 test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/launch.cpp create mode 100644 test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/main.cpp create mode 100644 test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/ptoas.flags create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/compare.py create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/golden.py create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/kernel.pto create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/launch.cpp create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/main.cpp create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/ptoas.flags create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/compare.py create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/golden.py create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/kernel.pto create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/launch.cpp create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/main.cpp create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/ptoas.flags create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/compare.py create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/golden.py create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/kernel.pto create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/launch.cpp create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/main.cpp create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/ptoas.flags create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/compare.py create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/golden.py create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/kernel.pto create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/launch.cpp create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/main.cpp create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/ptoas.flags create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/compare.py create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/golden.py create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/kernel.pto create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/launch.cpp create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/main.cpp create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/ptoas.flags create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/compare.py create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/golden.py create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/kernel.pto create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/launch.cpp create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/main.cpp create mode 100644 test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/ptoas.flags diff --git a/include/PTO/IR/PTOTypeDefs.td b/include/PTO/IR/PTOTypeDefs.td index 3a1c5022eb..569530185b 100644 --- a/include/PTO/IR/PTOTypeDefs.td +++ b/include/PTO/IR/PTOTypeDefs.td @@ -414,5 +414,14 @@ def F4E2M1x2Type : TypeDef +]> { + let mnemonic = "bf16x2"; + let summary = "Packed pair of BF16 values (4 bytes)."; +} + include "PTO/IR/VMITypeDefs.td" include "PTO/IR/VPTOTypeDefs.td" diff --git a/include/PTO/IR/PTOTypeUtils.h b/include/PTO/IR/PTOTypeUtils.h index 67fcf8b3a2..6cd1d7e671 100644 --- a/include/PTO/IR/PTOTypeUtils.h +++ b/include/PTO/IR/PTOTypeUtils.h @@ -42,6 +42,7 @@ bool isPTOFloat8E5M2LikeType(Type t); bool isPTOHiFloat8Type(Type t); bool isPTOF8E8M0Type(Type t); bool isPTOHiFloat8x2Type(Type t); +bool isPTOBF16x2Type(Type t); bool isPTOFloat4PackedType(Type t); bool isPTOPackedLdgStgVectorType(Type t); bool isPTOLowPrecisionType(Type t); diff --git a/include/PTO/IR/VMIUtils.h b/include/PTO/IR/VMIUtils.h index 1d2dc1b7a5..91ba18eb55 100644 --- a/include/PTO/IR/VMIUtils.h +++ b/include/PTO/IR/VMIUtils.h @@ -60,8 +60,8 @@ FailureOr isPaddingLane(Type type, int64_t part, int64_t chunk, // --------------------------------------------------------------------------- struct VMIFpToSiContract { - bool requiresSat; - bool requiresPart; + bool requiresSat = false; + bool requiresPart = false; }; /// Returns the FpToSi contract for the given src→dst element type pair, @@ -75,8 +75,8 @@ lookupVMIFpToSiContract(Type srcElem, Type dstElem); // --------------------------------------------------------------------------- struct VMIFpToUiContract { - bool requiresSat; - bool requiresPart; + bool requiresSat = false; + bool requiresPart = false; }; /// Returns the FpToUi contract for the given src→dst element type pair, @@ -86,14 +86,15 @@ lookupVMIFpToUIContract(Type srcElem, Type dstElem); // --------------------------------------------------------------------------- // VMI FpToFp hardware contract (VMI-owned; may diverge from VPTO). -// Only same-width fp->fp needs a pair whitelist here; widen/narrow fp->fp -// reuse the extf/truncf layout framework and do not consult this table. +// Enumerates same-width fp->fp whitelist pairs plus the fp->fp narrow paths +// whose sat semantics differ from the truncf default (e.g. bf16x2->f4x2). // --------------------------------------------------------------------------- struct VMIFpToFpContract { - bool requiresRnd; - bool requiresSat; - bool requiresPart; + bool requiresRnd = false; + bool requiresSat = false; + bool requiresPart = false; + StringRef allowedRndModes = StringRef(); }; /// Returns the FpToFp contract for the given src->dst element type pair, diff --git a/include/pto-c/Dialect/PTO.h b/include/pto-c/Dialect/PTO.h index cb653c2384..629cfda4af 100644 --- a/include/pto-c/Dialect/PTO.h +++ b/include/pto-c/Dialect/PTO.h @@ -50,6 +50,8 @@ bool mlirPTOTypeIsAF4E1M2x2Type(MlirType type); MlirType mlirPTOF4E1M2x2TypeGet(MlirContext ctx); bool mlirPTOTypeIsAF4E2M1x2Type(MlirType type); MlirType mlirPTOF4E2M1x2TypeGet(MlirContext ctx); +bool mlirPTOTypeIsABF16x2Type(MlirType type); +MlirType mlirPTOBF16x2TypeGet(MlirContext ctx); // ---- #pto.address_space<...> ---- bool mlirPTOAttrIsAAddressSpaceAttr(MlirAttribute attr); diff --git a/lib/Bindings/Python/PTOModule.cpp b/lib/Bindings/Python/PTOModule.cpp index a80981da98..2725d2ec70 100644 --- a/lib/Bindings/Python/PTOModule.cpp +++ b/lib/Bindings/Python/PTOModule.cpp @@ -1169,6 +1169,17 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { }, py::arg("cls"), py::arg("context") = py::none()); + mlir_type_subclass( + m, "BF16x2Type", + [](MlirType type) -> bool { return mlirPTOTypeIsABF16x2Type(type); }) + .def_classmethod( + "get", + [](py::object cls, MlirContext context) -> py::object { + MlirType t = mlirPTOBF16x2TypeGet(context); + return cls.attr("__call__")(t); + }, + py::arg("cls"), py::arg("context") = py::none()); + mlir_type_subclass( m, "F4E1M2x2Type", [](MlirType type) -> bool { return mlirPTOTypeIsAF4E1M2x2Type(type); }) diff --git a/lib/CAPI/Dialect/PTO.cpp b/lib/CAPI/Dialect/PTO.cpp index 61a37c11fa..0fc2b4d892 100644 --- a/lib/CAPI/Dialect/PTO.cpp +++ b/lib/CAPI/Dialect/PTO.cpp @@ -154,6 +154,14 @@ MlirType mlirPTOF4E2M1x2TypeGet(MlirContext ctx) { return wrap(mlir::pto::F4E2M1x2Type::get(unwrap(ctx))); } +bool mlirPTOTypeIsABF16x2Type(MlirType type) { + return isa(unwrap(type)); +} + +MlirType mlirPTOBF16x2TypeGet(MlirContext ctx) { + return wrap(mlir::pto::BF16x2Type::get(unwrap(ctx))); +} + MlirAttribute mlirPTOPtrTypeGetMemorySpace(MlirType type) { auto t = cast(unwrap(type)); return wrap(t.getMemorySpace()); diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index a33e08861d..427bb40d23 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -560,6 +560,25 @@ uint64_t mlir::pto::F4E2M1x2Type::getPreferredAlignment( return 1; } +static llvm::TypeSize getFourByteTypeSize() { + return llvm::TypeSize::getFixed(32); +} + +llvm::TypeSize mlir::pto::BF16x2Type::getTypeSizeInBits( + const DataLayout &, DataLayoutEntryListRef) const { + return getFourByteTypeSize(); +} + +uint64_t mlir::pto::BF16x2Type::getABIAlignment( + const DataLayout &, DataLayoutEntryListRef) const { + return 4; +} + +uint64_t mlir::pto::BF16x2Type::getPreferredAlignment( + const DataLayout &, DataLayoutEntryListRef) const { + return 4; +} + static VerifierTargetArch getVerifierTargetArch(Operation *op) { auto module = op ? op->getParentOfType() : ModuleOp(); if (isA5ModuleTarget(module)) diff --git a/lib/PTO/IR/PTOTypeUtils.cpp b/lib/PTO/IR/PTOTypeUtils.cpp index dc65a30834..9a70e4d12a 100644 --- a/lib/PTO/IR/PTOTypeUtils.cpp +++ b/lib/PTO/IR/PTOTypeUtils.cpp @@ -36,6 +36,8 @@ bool mlir::pto::isPTOF8E8M0Type(Type t) { return isa(t); } bool mlir::pto::isPTOHiFloat8x2Type(Type t) { return isa(t); } +bool mlir::pto::isPTOBF16x2Type(Type t) { return isa(t); } + bool mlir::pto::isPTOFloat4PackedType(Type t) { return isa(t); } @@ -82,13 +84,19 @@ unsigned mlir::pto::getPTOPackedLdgStgTotalBits(Type t) { bool mlir::pto::isPTOLowPrecisionType(Type t) { return isPTOFloat8Type(t) || isPTOHiFloat8Type(t) || isPTOF8E8M0Type(t) || - isPTOHiFloat8x2Type(t) || isPTOFloat4PackedType(t); + isPTOHiFloat8x2Type(t) || isPTOFloat4PackedType(t) || + isPTOBF16x2Type(t); } unsigned mlir::pto::getPTOStorageElemBitWidth(Type t) { if (isPTOHiFloat8x2Type(t)) { return 16; } + // bf16x2 is a 4-byte packed pair; special-case it before the generic + // low-precision branch (which would otherwise report 8 bits). + if (isPTOBF16x2Type(t)) { + return 32; + } if (isPTOLowPrecisionType(t)) return kBitsPerByte; if (auto floatTy = dyn_cast(t)) diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp index bc9697c2e1..7e761c5afa 100644 --- a/lib/PTO/IR/VMI.cpp +++ b/lib/PTO/IR/VMI.cpp @@ -59,6 +59,29 @@ static bool isVMIFloatLikeType(Type type) { return isa(type) || pto::isPTOLowPrecisionType(type); } +static bool involvesBF16x2(Type sourceType, Type resultType) { + return pto::isPTOBF16x2Type(sourceType) || + pto::isPTOBF16x2Type(resultType); +} + +static bool isVMIPackedFloatCarrierType(Type type) { + return pto::isPTOHiFloat8x2Type(type) || + pto::isPTOFloat4PackedType(type) || + pto::isPTOBF16x2Type(type); +} + +static bool involvesVMIPackedFloatCarrier(Type sourceType, Type resultType) { + return isVMIPackedFloatCarrierType(sourceType) || + isVMIPackedFloatCarrierType(resultType); +} + +static LogicalResult verifyBF16x2ComputeElementType(Operation *op, Type type) { + if (pto::isPTOBF16x2Type(type)) + return op->emitOpError( + "does not support bf16x2 VMI element type; bf16x2 is conversion-only"); + return success(); +} + static bool isVMIIntegerLikeType(Type type) { return isa(type); } @@ -171,16 +194,6 @@ static unsigned getVMIElementBitWidth(Type type) { return pto::getPTOStorageElemBitWidth(type); } -static std::optional getVMIIntegerOrFloatBitWidth(Type type) { - if (auto intType = dyn_cast(type)) { - return intType.getWidth(); - } - if (auto floatType = dyn_cast(type)) { - return floatType.getWidth(); - } - return std::nullopt; -} - static int64_t divideCeilNonNegative(int64_t value, int64_t divisor) { return value == 0 ? 0 : (value + divisor - 1) / divisor; } @@ -379,6 +392,10 @@ static LogicalResult verifyElementwiseVRegOp(Operation *op, VMIVRegType lhs, static LogicalResult verifyFloatUnaryVRegOp(Operation *op, VMIVRegType source, VMIVRegType result) { + if (failed( + verifyBF16x2ComputeElementType(op, source.getElementType()))) { + return failure(); + } if (!isVMIFloatLikeType(source.getElementType())) { return op->emitOpError("requires floating-point-like VMI element type"); } @@ -389,6 +406,9 @@ static LogicalResult verifyFloatUnaryVRegOp(Operation *op, VMIVRegType source, static LogicalResult verifyFloatTernaryVRegOp(Operation *op, VMIVRegType lhs, VMIVRegType rhs, VMIVRegType acc, VMIVRegType result) { + if (failed(verifyBF16x2ComputeElementType(op, lhs.getElementType()))) { + return failure(); + } if (!isVMIFloatLikeType(lhs.getElementType())) { return op->emitOpError("requires floating-point-like VMI element type"); } @@ -713,7 +733,9 @@ lookupVMIFpToUIContract(Type srcElem, Type dstElem) { // --------------------------------------------------------------------------- // FpToFp hardware contract (VMI-owned; may diverge from VPTO). -// Only same-width fp->fp is enumerated here. +// Enumerates same-width fp->fp whitelist entries plus the fp->fp narrow +// paths whose sat/rounding semantics differ from the generic truncf default +// (e.g. bf16x2->f4x2 narrows with NO saturation). // --------------------------------------------------------------------------- std::optional @@ -723,13 +745,29 @@ lookupVMIFpToFpContract(Type srcElem, Type dstElem) { } unsigned srcBits = pto::getPTOStorageElemBitWidth(srcElem); unsigned dstBits = pto::getPTOStorageElemBitWidth(dstElem); + // bf16x2 -> f4x2 (32->8 narrow): Packed4, rnd, NO sat. Mirrors the VPTO + // bf16->f4 contract row (requiresSat=false) so the VMI verifier does not + // force a saturate attribute that the physical pto.vcvt would reject. + if (pto::isPTOBF16x2Type(srcElem) && pto::isPTOFloat4PackedType(dstElem)) + return VMIFpToFpContract{/*requiresRnd=*/true, /*requiresSat=*/false, + /*requiresPart=*/true, + /*allowedRndModes=*/"RAFZC"}; + // f4x2 -> bf16x2 (8->32 widen): Packed4, no rnd, no sat. Mirrors the VPTO + // f4->bf16 contract row (requiresSat=false, requiresRnd=false). Widen has + // no rounding/saturate semantics by construction; the contract exists so + // the involvesBF16x2 / packed-carrier gates in the VMI verifiers pass. + if (pto::isPTOFloat4PackedType(srcElem) && pto::isPTOBF16x2Type(dstElem)) + return VMIFpToFpContract{/*requiresRnd=*/false, /*requiresSat=*/false, + /*requiresPart=*/true, + /*allowedRndModes=*/StringRef()}; if (srcBits != dstBits) { return std::nullopt; } // bf16 -> f16: same-width, rnd, sat, no part. if (srcElem.isBF16() && dstElem.isF16()) return VMIFpToFpContract{/*requiresRnd=*/true, /*requiresSat=*/true, - /*requiresPart=*/false}; + /*requiresPart=*/false, + /*allowedRndModes=*/StringRef()}; return std::nullopt; } @@ -976,14 +1014,6 @@ LogicalResult VMIVRegType::verify(function_ref emitError, << formatVMIVRegType(elementCount, elementType, layout) << "' expected an 8-bit, 16-bit, or 32-bit logical " "element type"; - if (pto::isPTOFloat4PackedType(elementType)) - return emitError() - << "'" << formatVMIVRegType(elementCount, elementType, layout) - << "' uses a packed FP4 physical pair type as a VMI logical " - "element type; packed FP4 input/output is not a supported VMI " - "surface because the logical FP4 lane count and physical packed " - "byte count are ambiguous"; - if (layout && !mlir::isa(layout)) return emitError() << "'" << formatVMIVRegType(elementCount, elementType, layout) @@ -1923,15 +1953,21 @@ LogicalResult VMIVchistOp::verify() { return verifyVMIHistogramOp(*this); } LogicalResult VMIExtFOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); + Type sourceElementType = sourceType.getElementType(); + Type resultElementType = resultType.getElementType(); if (sourceType.getElementCount() != resultType.getElementCount()) return emitOpError( "requires source and result logical lane counts to match"); - if (!isVMIFloatLikeType(sourceType.getElementType()) || - !isVMIFloatLikeType(resultType.getElementType())) + if (!isVMIFloatLikeType(sourceElementType) || + !isVMIFloatLikeType(resultElementType)) return emitOpError( "requires floating-point-like source and result element types"); - if (getVMIElementBitWidth(sourceType.getElementType()) >= - getVMIElementBitWidth(resultType.getElementType())) + if (involvesBF16x2(sourceElementType, resultElementType) && + !lookupVMIFpToFpContract(sourceElementType, resultElementType)) + return emitOpError( + "unsupported bf16x2 fp-to-fp conversion element type pair"); + if (getVMIElementBitWidth(sourceElementType) >= + getVMIElementBitWidth(resultElementType)) return emitOpError( "requires result element type to be wider than source element type"); return success(); @@ -1940,15 +1976,26 @@ LogicalResult VMIExtFOp::verify() { LogicalResult VMITruncFOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); + Type sourceElementType = sourceType.getElementType(); + Type resultElementType = resultType.getElementType(); + auto fpContract = + lookupVMIFpToFpContract(sourceElementType, resultElementType); if (sourceType.getElementCount() != resultType.getElementCount()) return emitOpError( "requires source and result logical lane counts to match"); - if (!isVMIFloatLikeType(sourceType.getElementType()) || - !isVMIFloatLikeType(resultType.getElementType())) + if (!isVMIFloatLikeType(sourceElementType) || + !isVMIFloatLikeType(resultElementType)) return emitOpError( "requires floating-point-like source and result element types"); - unsigned srcBits = getVMIElementBitWidth(sourceType.getElementType()); - unsigned dstBits = getVMIElementBitWidth(resultType.getElementType()); + if (involvesBF16x2(sourceElementType, resultElementType) && !fpContract) + return emitOpError( + "unsupported bf16x2 fp-to-fp conversion element type pair"); + if (involvesVMIPackedFloatCarrier(sourceElementType, resultElementType) && + !fpContract) + return emitOpError( + "unsupported packed fp-to-fp conversion element type pair"); + unsigned srcBits = getVMIElementBitWidth(sourceElementType); + unsigned dstBits = getVMIElementBitWidth(resultElementType); if (srcBits < dstBits) return emitOpError( "requires result element type to be narrower than or same-width " @@ -1956,24 +2003,42 @@ LogicalResult VMITruncFOp::verify() { if (srcBits == dstBits) { // Same-width fp→fp (e.g. bf16→f16): only allowed for supported VMI // fp-to-fp contract pairs. - if (!lookupVMIFpToFpContract(sourceType.getElementType(), - resultType.getElementType())) + if (!fpContract) return emitOpError("same-width fp-to-fp conversion is not supported " "for this type pair; see lookupVMIFpToFpContract"); } if (auto roundingAttr = (*this)->getAttrOfType("rounding")) { StringRef rounding = roundingAttr.getValue(); - if (rounding != "R" && rounding != "A" && rounding != "H" && - rounding != "Z") + if (rounding.size() != 1) + return emitOpError( + "rounding attr must be a single-character mode token"); + StringRef allowedRndModes = + fpContract && !fpContract->allowedRndModes.empty() + ? fpContract->allowedRndModes + : StringRef("RAHZ"); + if (!allowedRndModes.contains(rounding)) { + if (fpContract && !fpContract->allowedRndModes.empty()) + return emitOpError("rounding attr is not valid for this fp-to-fp " + "conversion type pair"); return emitOpError("rounding attr must be R, A, H, or Z"); + } } auto satAttr = (*this)->getAttrOfType("saturate"); - if (!satAttr) { - return emitOpError("'saturate' attribute is required (SAT or NOSAT)"); - } - StringRef satVal = satAttr.getValue(); - if (satVal != "SAT" && satVal != "NOSAT") { - return emitOpError("saturate attr must be 'SAT' or 'NOSAT'"); + // Some fp->fp narrow paths (e.g. bf16x2 -> f4x2) do NOT saturate; consult + // the fp-to-fp contract when one exists instead of always requiring SAT. + if (!fpContract || fpContract->requiresSat) { + if (!satAttr) { + return emitOpError("'saturate' attribute is required (SAT or NOSAT)"); + } + StringRef satVal = satAttr.getValue(); + if (satVal != "SAT" && satVal != "NOSAT") { + return emitOpError("saturate attr must be 'SAT' or 'NOSAT'"); + } + } else { + if (satAttr) { + return emitOpError("'saturate' attribute is not valid for this fp-to-fp " + "narrow conversion (no saturation)"); + } } return success(); } @@ -2130,15 +2195,15 @@ LogicalResult VMITruncIOp::verify() { LogicalResult VMIBitcastOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - std::optional sourceBits = - getVMIIntegerOrFloatBitWidth(sourceType.getElementType()); - std::optional resultBits = - getVMIIntegerOrFloatBitWidth(resultType.getElementType()); - if (!sourceBits || !resultBits) + unsigned sourceBits = + pto::getPTOStorageElemBitWidth(sourceType.getElementType()); + unsigned resultBits = + pto::getPTOStorageElemBitWidth(resultType.getElementType()); + if (sourceBits == 0 || resultBits == 0) return emitOpError( "requires integer or floating-point source and result element types"); - if (sourceType.getElementCount() * static_cast(*sourceBits) != - resultType.getElementCount() * static_cast(*resultBits)) + if (sourceType.getElementCount() * static_cast(sourceBits) != + resultType.getElementCount() * static_cast(resultBits)) return emitOpError( "requires source and result to carry the same total number of bits"); @@ -2736,6 +2801,8 @@ verifyVMIVectorScalarOp(Operation *op, VMIVRegType srcType, VMIMaskType maskType, std::optional pmode) { Type eltTy = srcType.getElementType(); + if (failed(verifyBF16x2ComputeElementType(op, eltTy))) + return failure(); if (!isVMIFloatLikeType(eltTy) && !isVMIIntegerLikeType(eltTy)) return op->emitOpError( "requires floating-point-like or integer-like VMI element type"); @@ -3002,6 +3069,10 @@ LogicalResult VMIVaddOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); + if (failed(verifyBF16x2ComputeElementType( + getOperation(), lhsType.getElementType()))) { + return failure(); + } if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) { return failure(); } @@ -3015,6 +3086,10 @@ LogicalResult VMIVsubOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); + if (failed(verifyBF16x2ComputeElementType( + getOperation(), lhsType.getElementType()))) { + return failure(); + } if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) { return failure(); } @@ -3028,6 +3103,10 @@ LogicalResult VMIVmulOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); + if (failed(verifyBF16x2ComputeElementType( + getOperation(), lhsType.getElementType()))) { + return failure(); + } if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) { return failure(); } @@ -3041,6 +3120,10 @@ LogicalResult VMIVdivOp::verify() { auto lhsType = cast(getLhs().getType()); auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); + if (failed(verifyBF16x2ComputeElementType( + getOperation(), lhsType.getElementType()))) { + return failure(); + } if (!isVMIFloatLikeType(lhsType.getElementType())) { return emitOpError("requires floating-point-like VMI element type"); } @@ -3058,6 +3141,8 @@ LogicalResult VMIVminOp::verify() { auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); Type elementType = lhsType.getElementType(); + if (failed(verifyBF16x2ComputeElementType(getOperation(), elementType))) + return failure(); if (!isVMIFloatLikeType(elementType) && !isVMIAnyI8I16I32Type(elementType)) return emitOpError( "requires floating-point-like or i8, i16, or i32 VMI element type"); @@ -3075,6 +3160,8 @@ LogicalResult VMIVmaxOp::verify() { auto rhsType = cast(getRhs().getType()); auto resultType = cast(getResult().getType()); Type elementType = lhsType.getElementType(); + if (failed(verifyBF16x2ComputeElementType(getOperation(), elementType))) + return failure(); if (!isVMIFloatLikeType(elementType) && !isVMIAnyI8I16I32Type(elementType)) return emitOpError( "requires floating-point-like or i8, i16, or i32 VMI element type"); @@ -3104,6 +3191,8 @@ LogicalResult VMIVabsOp::verify() { auto resultType = cast(getResult().getType()); Type eltTy = sourceType.getElementType(); + if (failed(verifyBF16x2ComputeElementType(getOperation(), eltTy))) + return failure(); if (!isVMIFloatLikeType(eltTy) && !isVMIIntegerLikeType(eltTy)) return emitOpError( "requires floating-point-like or integer-like VMI element type"); @@ -3350,6 +3439,9 @@ LogicalResult VMIvcaddOp::verify() { auto resultType = cast(getResult().getType()); auto elemTy = sourceType.getElementType(); + if (failed(verifyBF16x2ComputeElementType(getOperation(), elemTy))) + return failure(); + // Element type must be integer-like or float-like bool isFloat = isVMIFloatLikeType(elemTy); bool isInt = isVMIIntegerLikeType(elemTy); @@ -3411,6 +3503,9 @@ LogicalResult VMIvcmaxOp::verify() { auto resultType = cast(getResult().getType()); auto elemTy = sourceType.getElementType(); + if (failed(verifyBF16x2ComputeElementType(getOperation(), elemTy))) + return failure(); + bool isFloat = isVMIFloatLikeType(elemTy); bool isInt = isVMIIntegerLikeType(elemTy); if (!isFloat && !isInt) @@ -3463,6 +3558,9 @@ LogicalResult VMIvcminOp::verify() { auto resultType = cast(getResult().getType()); auto elemTy = sourceType.getElementType(); + if (failed(verifyBF16x2ComputeElementType(getOperation(), elemTy))) + return failure(); + bool isFloat = isVMIFloatLikeType(elemTy); bool isInt = isVMIIntegerLikeType(elemTy); if (!isFloat && !isInt) @@ -3650,6 +3748,11 @@ LogicalResult VMIVexpdifOp::verify() { auto maskType = cast(getMask().getType()); auto resultType = cast(getResult().getType()); + if (failed(verifyBF16x2ComputeElementType( + getOperation(), xType.getElementType()))) { + return failure(); + } + if (!isVMIFloatLikeType(xType.getElementType())) { return emitOpError("requires x element type to be f16 or f32"); } @@ -3687,6 +3790,10 @@ LogicalResult VMIVaxpyOp::verify() { auto maskType = cast(getMask().getType()); auto resultType = cast(getResult().getType()); + if (failed(verifyBF16x2ComputeElementType( + getOperation(), xType.getElementType()))) { + return failure(); + } if (!isVMIFloatLikeType(xType.getElementType())) { return emitOpError("requires vector element type to be f16 or f32"); } @@ -3717,6 +3824,10 @@ LogicalResult VMIVlreluOp::verify() { auto maskType = cast(getMask().getType()); auto resultType = cast(getResult().getType()); + if (failed(verifyBF16x2ComputeElementType( + getOperation(), xType.getElementType()))) { + return failure(); + } if (!isVMIFloatLikeType(xType.getElementType())) { return emitOpError("requires vector element type to be f16 or f32"); } @@ -3748,6 +3859,11 @@ LogicalResult VMIVpreluOp::verify() { auto maskType = cast(getMask().getType()); auto resultType = cast(getResult().getType()); + if (failed(verifyBF16x2ComputeElementType( + getOperation(), xType.getElementType()))) { + return failure(); + } + if (!isVMIFloatLikeType(xType.getElementType())) { return emitOpError("requires vector element type to be f16 or f32"); } @@ -3852,6 +3968,8 @@ LogicalResult VMIVmulaOp::verify() { auto resultType = cast(getResult().getType()); Type eltTy = accType.getElementType(); + if (failed(verifyBF16x2ComputeElementType(getOperation(), eltTy))) + return failure(); if (!isVMIFloatLikeType(eltTy) && !isVMIIntegerLikeType(eltTy)) return emitOpError( "requires floating-point-like or integer-like VMI element type"); @@ -3887,6 +4005,17 @@ LogicalResult VMICvtOp::verify() { srcInt && isa(srcElem) && !cast(srcElem).isUnsigned() && !cast(srcElem).isSigned(); + auto fpContract = + srcFp && dstFp ? lookupVMIFpToFpContract(srcElem, dstElem) + : std::nullopt; + + if (involvesBF16x2(srcElem, dstElem) && !fpContract) + return emitOpError( + "unsupported conversion involving bf16x2 element type"); + if (srcFp && dstFp && + involvesVMIPackedFloatCarrier(srcElem, dstElem) && !fpContract) + return emitOpError( + "unsupported packed fp-to-fp conversion element type pair"); // 2. Classify the conversion direction. CvtDirection dir; @@ -3900,7 +4029,7 @@ LogicalResult VMICvtOp::verify() { else { // Same-width fp→fp (e.g. bf16 → f16): only allowed for VMI fp-to-fp // contract pairs, routed through FpNarrow (1:1 TruncF). - if (!lookupVMIFpToFpContract(srcElem, dstElem)) + if (!fpContract) return emitOpError( "same-width fp-to-fp conversion is not supported for this type " "pair; see lookupVMIFpToFpContract"); @@ -3942,10 +4071,21 @@ LogicalResult VMICvtOp::verify() { return emitOpError("'rounding' attribute is only valid for " "fp-narrowing conversions"); StringRef rnd = roundingAttr.getValue(); - if (rnd != "R" && rnd != "A" && rnd != "H" && rnd != "Z") + if (rnd.size() != 1) + return emitOpError( + "rounding must be a single-character mode token"); + StringRef allowedRndModes = + fpContract && !fpContract->allowedRndModes.empty() + ? fpContract->allowedRndModes + : StringRef("RAHZ"); + if (!allowedRndModes.contains(rnd)) { + if (fpContract && !fpContract->allowedRndModes.empty()) + return emitOpError( + "rounding is not valid for this fp-to-fp conversion type pair"); return emitOpError("rounding must be 'R' (nearest-even), " "'A' (away-from-zero), 'H' (half-up), " "or 'Z' (toward-zero)"); + } } // --- saturate --- @@ -3987,8 +4127,13 @@ LogicalResult VMICvtOp::verify() { "fp-to-ui conversion (no overflow possible)"); } } else { - bool needSat = (dir == CvtDirection::FpNarrow || - dir == CvtDirection::IntNarrow); + bool needSat = (dir == CvtDirection::IntNarrow); + // Fp-narrow: default to requiring a saturate attribute, but consult the + // fp-to-fp contract when one exists (e.g. bf16x2->f4x2 narrows with + // requiresSat=false and must NOT carry saturate). + if (dir == CvtDirection::FpNarrow) { + needSat = !fpContract || fpContract->requiresSat; + } if (needSat) { if (!satAttr) return emitOpError("'saturate' attribute is required for fp-narrow / " @@ -4011,6 +4156,9 @@ LogicalResult VMICvtOp::verify() { return emitOpError("si32 -> si8 int-narrow does not support " "saturate=\"SAT\" (no native hardware form; " "only saturate=\"NOSAT\" is allowed)"); + } else if (satAttr && dir == CvtDirection::FpNarrow) { + return emitOpError("'saturate' attribute is not valid for this fp-to-fp " + "narrow conversion (no saturation)"); } else if (satAttr) { return emitOpError("'saturate' attribute is only valid for fp-narrow / " "int-narrow conversions"); @@ -4041,15 +4189,15 @@ LogicalResult VMICvtOp::verify() { LogicalResult VMIVinterpretCastOp::verify() { auto sourceType = cast(getSource().getType()); auto resultType = cast(getResult().getType()); - std::optional sourceBits = - getVMIIntegerOrFloatBitWidth(sourceType.getElementType()); - std::optional resultBits = - getVMIIntegerOrFloatBitWidth(resultType.getElementType()); - if (!sourceBits || !resultBits) + unsigned sourceBits = + pto::getPTOStorageElemBitWidth(sourceType.getElementType()); + unsigned resultBits = + pto::getPTOStorageElemBitWidth(resultType.getElementType()); + if (sourceBits == 0 || resultBits == 0) return emitOpError( "requires integer or floating-point source and result element types"); - if (sourceType.getElementCount() * static_cast(*sourceBits) != - resultType.getElementCount() * static_cast(*resultBits)) + if (sourceType.getElementCount() * static_cast(sourceBits) != + resultType.getElementCount() * static_cast(resultBits)) return emitOpError( "requires source and result to carry the same total number of bits"); @@ -4484,6 +4632,8 @@ LogicalResult VMIVcmpOp::verify() { // Element type must be float-like OR integer-like (unified). Type eltTy = lhsType.getElementType(); + if (failed(verifyBF16x2ComputeElementType(getOperation(), eltTy))) + return failure(); if (!isVMIFloatLikeType(eltTy) && !isVMIIntegerLikeType(eltTy)) return emitOpError("requires floating-point-like or integer-like VMI " "element type for unified compare"); @@ -4528,6 +4678,8 @@ LogicalResult VMIVcmpsOp::verify() { // Element type must be float-like OR integer-like (unified). Type eltTy = srcType.getElementType(); + if (failed(verifyBF16x2ComputeElementType(getOperation(), eltTy))) + return failure(); if (!isVMIFloatLikeType(eltTy) && !isVMIIntegerLikeType(eltTy)) return emitOpError("requires floating-point-like or integer-like VMI " "element type for unified compare"); diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index a63c4bc468..67f4630c73 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -6044,6 +6044,13 @@ static LogicalResult verifyBinaryVecOp(BinaryOp op, failed(verifyNonLowPrecisionVRegElementTypeLike( op.getOperation(), op.getLhs().getType(), "lhs type"))) return failure(); + if (allowLowPrecision) { + auto lhsType = cast(op.getLhs().getType()); + if (pto::isPTOBF16x2Type(lhsType.getElementType())) + return op.emitOpError( + "does not support bf16x2 vector elements; low-precision bitwise " + "operations require an 8-bit payload type"); + } if (op.getLhs().getType() != op.getRhs().getType() || op.getLhs().getType() != op.getResult().getType()) return op.emitOpError("requires matching register vector shapes"); @@ -6702,6 +6709,11 @@ LogicalResult VbitcastOp::verify() { if (auto floatType = dyn_cast(elementType)) return type.getElementCount() * static_cast(floatType.getWidth()); + // Packed PTO element types (f8/hif8/f4x2/bf16x2/...) have a known storage + // width even though they are not IntegerType/FloatType. + unsigned packedBits = pto::getPTOStorageElemBitWidth(elementType); + if (packedBits != 0) + return type.getElementCount() * static_cast(packedBits); return std::nullopt; }; diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 52253bd2ec..f4c0e6bdb7 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -157,6 +157,12 @@ bool hasVMIType(Operation *op) { return false; } +bool isVMIPackedFloatCarrierType(Type type) { + return pto::isPTOHiFloat8x2Type(type) || + pto::isPTOFloat4PackedType(type) || + pto::isPTOBF16x2Type(type); +} + bool isVMIOp(Operation *op) { return op->getName().getStringRef().starts_with("pto.vmi."); } @@ -11093,8 +11099,11 @@ struct OneToNVMIExtFOpPattern : OpConversionPattern { for (Type resultType : resultTypes) { auto resultVRegType = dyn_cast(resultType); if (!resultVRegType || - (resultVRegTypes.empty() ? !resultVRegType.getElementType().isF32() - : resultVRegType != resultVRegTypes.front())) + (resultVRegTypes.empty() + ? !(resultVRegType.getElementType().isF32() || + pto::isPTOBF16x2Type( + resultVRegType.getElementType())) + : resultVRegType != resultVRegTypes.front())) return rewriter.notifyMatchFailure( op, "unsupported physical extf result type"); resultVRegTypes.push_back(resultVRegType); @@ -11102,6 +11111,36 @@ struct OneToNVMIExtFOpPattern : OpConversionPattern { unsigned sourceBits = pto::getPTOStorageElemBitWidth(sourceType.getElementType()); + // A packed bf16x2 physical result cannot be produced directly by + // pto.vcvt (classifyVcvtElemType has no BF16x2 branch); the widest native + // f4 conversion result element is bf16. Build the bf16 view type (2 bf16 + // lanes per bf16x2 lane) and reinterpret each vcvt result with a + // physical-noop VbitcastOp, mirroring the source-side reinterpret in + // OneToNVMITruncFOpPattern (viewVcvtSource). + bool resultIsPackedBF16x2 = + pto::isPTOBF16x2Type(resultVRegTypes.front().getElementType()); + VRegType vcvtResultVRegType = resultVRegTypes.front(); + if (resultIsPackedBF16x2) { + vcvtResultVRegType = + VRegType::get(rewriter.getContext(), + resultVRegTypes.front().getElementCount() * 2, + BFloat16Type::get(rewriter.getContext())); + } + auto viewVcvtResult = [&](VRegType resultType, Value sourcePart, + Value mask, StringAttr rnd, StringAttr sat, + StringAttr part) -> Value { + VRegType vcvtType = + resultIsPackedBF16x2 ? vcvtResultVRegType : resultType; + Value vcvt = rewriter + .create(op.getLoc(), vcvtType, sourcePart, mask, + rnd, sat, part) + .getResult(); + if (!resultIsPackedBF16x2) + return vcvt; + return rewriter.create(op.getLoc(), resultType, vcvt) + .getResult(); + }; + VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); if (sourceLayout && resultLayout && sourceLayout.isContiguous() && @@ -11120,12 +11159,9 @@ struct OneToNVMIExtFOpPattern : OpConversionPattern { results.reserve(resultTypes.size()); for (auto [sourcePart, resultType] : llvm::zip_equal(sourceParts, resultVRegTypes)) { - results.push_back(rewriter - .create(op.getLoc(), resultType, - sourcePart, *mask, - /*rnd=*/nullptr, /*sat=*/nullptr, - rewriter.getStringAttr(part)) - .getResult()); + results.push_back(viewVcvtResult( + resultType, sourcePart, *mask, /*rnd=*/nullptr, /*sat=*/nullptr, + rewriter.getStringAttr(part))); } replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); return success(); @@ -11159,12 +11195,9 @@ struct OneToNVMIExtFOpPattern : OpConversionPattern { for (auto [chunkIndex, sourcePart] : llvm::enumerate(sourceParts)) { VRegType resultType = resultVRegTypes[partIndex * sourceParts.size() + chunkIndex]; - results.push_back( - rewriter - .create(op.getLoc(), resultType, sourcePart, *mask, - /*rnd=*/nullptr, /*sat=*/nullptr, - rewriter.getStringAttr(parts[partIndex])) - .getResult()); + results.push_back(viewVcvtResult( + resultType, sourcePart, *mask, /*rnd=*/nullptr, /*sat=*/nullptr, + rewriter.getStringAttr(parts[partIndex]))); } } @@ -11181,6 +11214,13 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { ConversionPatternRewriter &rewriter) const override { auto sourceVMIType = cast(op.getSource().getType()); auto resultVMIType = cast(op.getResult().getType()); + Type sourceElementType = sourceVMIType.getElementType(); + Type resultElementType = resultVMIType.getElementType(); + if ((isVMIPackedFloatCarrierType(sourceElementType) || + isVMIPackedFloatCarrierType(resultElementType)) && + !lookupVMIFpToFpContract(sourceElementType, resultElementType)) + return rewriter.notifyMatchFailure( + op, "unsupported packed fp-to-fp truncf conversion"); ValueRange sourceParts = adaptor.getSource(); FailureOr> maybe_resultTypes = getConvertedResultTypes(op, 0, *this->getTypeConverter()); @@ -11251,13 +11291,42 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { } auto sourceType0 = dyn_cast(sourceParts.front().getType()); - if (!sourceType0 || !isa(sourceType0.getElementType())) { + if (!sourceType0) { return rewriter.notifyMatchFailure(op, "unsupported physical truncf source type"); } unsigned sourceBits = pto::getPTOStorageElemBitWidth(sourceType0.getElementType()); if (sourceBits != 32 && sourceBits != 16) return rewriter.notifyMatchFailure( op, "truncf source bit width must be 32 or 16"); + // A packed bf16x2 physical source is consumed by pto.vcvt as raw bf16 + // lanes (2 bf16 per bf16x2). Build the bf16 view type used for the source + // mask and for reinterpreting each source part before the VcvtOp. The + // logical lane count stays bf16x2-based; only the physical view widens. + bool sourceIsPackedBF16x2 = + pto::isPTOBF16x2Type(sourceType0.getElementType()); + VRegType vcvtSourceVRegType = sourceType0; + if (sourceIsPackedBF16x2) { + vcvtSourceVRegType = + VRegType::get(rewriter.getContext(), sourceType0.getElementCount() * 2, + BFloat16Type::get(rewriter.getContext())); + } + auto viewVcvtSource = [&](Value sourcePart) -> Value { + if (!sourceIsPackedBF16x2) + return sourcePart; + // If the source part is the physical noop pairing bitcast produced by + // VMIBitcastOp lowering (bf16 vreg -> bf16x2 vreg), reuse the original + // bf16 value directly instead of re-viewing it. This avoids emitting a + // redundant view vbitcast and leaves the pairing bitcast dead so the + // emitter can erase it. + if (auto vbc = sourcePart.getDefiningOp()) { + if (auto srcVReg = dyn_cast(vbc.getInput().getType()); + srcVReg && srcVReg.getElementType().isBF16()) + return vbc.getInput(); + } + return rewriter + .create(op.getLoc(), vcvtSourceVRegType, sourcePart) + .getResult(); + }; // Group-slot layout for non-f32 sources is not supported yet. if (sourceLayout && sourceLayout.isGroupSlots()) return rewriter.notifyMatchFailure( @@ -11290,7 +11359,7 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { resultLayout.isContiguous() && resultLayout.getLaneStride() == 1 && sourceParts.size() == resultTypes.size()) { FailureOr sourceMask = - createAllTrueMaskForVReg(op.getLoc(), sourceType0, rewriter); + createAllTrueMaskForVReg(op.getLoc(), vcvtSourceVRegType, rewriter); if (failed(sourceMask)) { return rewriter.notifyMatchFailure(op, "failed to build truncf masks"); } @@ -11329,7 +11398,7 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { op, "unsupported dense lane_stride truncf result layout"); FailureOr sourceMask = - createAllTrueMaskForVReg(op.getLoc(), sourceType0, rewriter); + createAllTrueMaskForVReg(op.getLoc(), vcvtSourceVRegType, rewriter); if (failed(sourceMask)) { return rewriter.notifyMatchFailure(op, "failed to build truncf masks"); } @@ -11344,8 +11413,8 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { llvm::zip_equal(sourceParts, resultVRegTypes)) { results.push_back(rewriter .create(op.getLoc(), resultType, - sourcePart, *sourceMask, rnd, sat, - partAttr) + viewVcvtSource(sourcePart), + *sourceMask, rnd, sat, partAttr) .getResult()); } replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); @@ -11381,7 +11450,7 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { op, "unsupported physical truncf source/result arity relation"); FailureOr sourceMask = - createAllTrueMaskForVReg(op.getLoc(), sourceType0, rewriter); + createAllTrueMaskForVReg(op.getLoc(), vcvtSourceVRegType, rewriter); if (failed(sourceMask)) { return rewriter.notifyMatchFailure(op, "failed to build truncf masks"); } @@ -11405,8 +11474,9 @@ struct OneToNVMITruncFOpPattern : OpConversionPattern { sourceParts[partIndex * resultTypes.size() + chunkIndex]; partials.push_back( rewriter - .create(op.getLoc(), resultType, sourcePart, - *sourceMask, rnd, sat, + .create(op.getLoc(), resultType, + viewVcvtSource(sourcePart), *sourceMask, rnd, + sat, rewriter.getStringAttr( allParts[partIndex * resultLaneStride])) .getResult()); diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index 53adc3d39c..a19e9f176a 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -96,6 +96,11 @@ static Type normalizePayloadTypeForLLVMLowering(Type type, Builder &builder) { if (pto::isPTOHiFloat8x2Type(type)) return getLLVMCompatibleVectorType( {2}, LLVM::LLVMHiFloat8Type::get(builder.getContext())); + // bf16x2 is a 4-byte packed pair; lower it as an opaque i32 so vregs whose + // element type is bf16x2 get a valid LLVM type. + if (pto::isPTOBF16x2Type(type)) { + return builder.getI32Type(); + } if (Type lowpType = getLowPrecisionLLVMType(type, builder.getContext())) { return lowpType; } @@ -125,6 +130,10 @@ static Type normalizeGEPElementTypeForLLVMLowering(Type type, if (pto::isPTOHiFloat8x2Type(type)) { return builder.getI16Type(); } + // bf16x2 is 4 bytes, not an 8-bit low-precision type. + if (pto::isPTOBF16x2Type(type)) { + return builder.getI32Type(); + } if (pto::isPTOLowPrecisionType(type)) { return builder.getI8Type(); } @@ -189,6 +198,9 @@ static unsigned getNaturalByteAlignment(Type type) { if (pto::isPTOHiFloat8x2Type(type)) { return 2; } + if (pto::isPTOBF16x2Type(type)) { + return 4; + } if (pto::isPTOLowPrecisionType(type)) { return 1; } @@ -669,6 +681,8 @@ static std::string getLowPrecisionElementFragment(Type type) { return "f4e1m2x2"; if (isa(type)) return "f4e2m1x2"; + if (pto::isPTOBF16x2Type(type)) + return "bf16x2"; if (pto::isPTOFloat8E4M3LikeType(type)) return "f8e4m3"; if (pto::isPTOFloat8E5M2LikeType(type)) @@ -1231,6 +1245,9 @@ static std::optional getDistElementWidth(Type type) { return 32; if (type.isF64()) return 64; + // bf16x2 is a 32-bit packed pair; its dist width is 32 (i32/align4 ABI). + if (pto::isPTOBF16x2Type(type)) + return 32; return std::nullopt; } @@ -3658,6 +3675,16 @@ StringRef buildPredicatePairReorderCallee(MLIRContext *context static FailureOr buildInterleaveCallee(MLIRContext *context, Type resultType, StringRef stem) { + // bf16x2 has no dedicated vintlv/vdintlv intrinsic. It is a 32-bit packed + // pair lowered to i32 at the LLVM ABI, and (de)interleave is a bit-level + // lane shuffle, so the intrinsic serves the type. + if (pto::isPTOBF16x2Type(getElementTypeFromVectorLike(resultType))) { + auto lanes = getElementCountFromVectorLike(resultType); + if (lanes) + return StringAttr::get(context, "llvm.hivm." + stem.str() + ".v" + + std::to_string(*lanes) + "i32") + .getValue(); + } std::string vec = getCANN900VectorTypeFragment(resultType); if (vec.empty()) return failure(); @@ -8266,6 +8293,13 @@ class LowerVbitcastOpPattern final LogicalResult matchAndRewrite(pto::VbitcastOp op, pto::VbitcastOp::Adaptor adaptor, ConversionPatternRewriter &rewriter) const override { + // A vbitcast whose result has no users is a dead noop (Pure). Erase it + // instead of emitting an LLVM bitcast the device compiler may not lower + // (e.g. bf16x2 <-> bf16 physical views). + if (op->use_empty()) { + rewriter.eraseOp(op); + return success(); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index d9cf37e156..3e5de0fd4d 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -103,6 +103,10 @@ static Type normalizePayloadTypeForLLVMLowering(Type type, Builder &builder) { if (pto::isPTOHiFloat8x2Type(type)) return getLLVMCompatibleVectorType( {2}, LLVM::LLVMHiFloat8Type::get(builder.getContext())); + // bf16x2 is a 4-byte packed pair; lower it as an opaque i32 so vregs whose + // element type is bf16x2 get a valid LLVM type. + if (pto::isPTOBF16x2Type(type)) + return builder.getI32Type(); if (Type lowpType = getLowPrecisionLLVMType(type, builder.getContext())) { return lowpType; @@ -136,6 +140,10 @@ static Type normalizeGEPElementTypeForLLVMLowering(Type type, { return builder.getI16Type(); } + // bf16x2 is 4 bytes, not an 8-bit low-precision type. + if (pto::isPTOBF16x2Type(type)) { + return builder.getI32Type(); + } if (pto::isPTOLowPrecisionType(type)) { return builder.getI8Type(); @@ -205,6 +213,9 @@ static unsigned getNaturalByteAlignment(Type type) { { return 2; } + if (pto::isPTOBF16x2Type(type)) { + return 4; + } if (pto::isPTOLowPrecisionType(type)) { return 1; @@ -690,6 +701,9 @@ static std::string getLowPrecisionElementFragment(Type type) { { return "f4e2m1x2"; } + if (pto::isPTOBF16x2Type(type)) { + return "bf16x2"; + } if (pto::isPTOFloat8E4M3LikeType(type)) { return "f8e4m3"; @@ -1493,6 +1507,10 @@ static std::optional getDistElementWidth(Type type) { { return 64; } + // bf16x2 is a 32-bit packed pair; its dist width is 32 (i32/align4 ABI). + if (pto::isPTOBF16x2Type(type)) { + return 32; + } return std::nullopt; } @@ -2762,6 +2780,8 @@ static FailureOr convertElementOffsetToBytes(Operation *anchor, Value off bitWidth = 8; else if (auto floatType = dyn_cast(elementType)) bitWidth = floatType.getWidth(); + else if (pto::isPTOBF16x2Type(elementType)) + bitWidth = 32; if (bitWidth == 0 || bitWidth % 8 != 0) { return failure(); @@ -4438,6 +4458,16 @@ StringRef buildPredicatePairReorderCallee(MLIRContext *context static FailureOr buildInterleaveCallee(MLIRContext *context, Type resultType, StringRef stem) { + // bf16x2 has no dedicated vintlv/vdintlv intrinsic: it is a 32-bit packed + // pair lowered to i32 at the LLVM ABI, and (de)interleave is a bit-level + // lane shuffle, so the intrinsic serves the type. + if (pto::isPTOBF16x2Type(getElementTypeFromVectorLike(resultType))) { + auto lanes = getElementCountFromVectorLike(resultType); + if (lanes) + return StringAttr::get(context, "llvm.hivm." + stem.str() + ".v" + + std::to_string(*lanes) + "i32") + .getValue(); + } return buildLaneTypedCallee(context, resultType, stem, ""); } @@ -10032,6 +10062,13 @@ class LowerVbitcastOpPattern final LogicalResult matchAndRewrite(pto::VbitcastOp op, pto::VbitcastOp::Adaptor adaptor, ConversionPatternRewriter &rewriter) const override { + // A vbitcast whose result has no users is a dead noop (Pure). Erase it + // instead of emitting an LLVM bitcast the device compiler may not lower + // (e.g. bf16x2 <-> bf16 physical views). + if (op->use_empty()) { + rewriter.eraseOp(op); + return success(); + } Type resultType = this->getTypeConverter()->convertType(op.getResult().getType()); if (!resultType) diff --git a/ptodsl/ptodsl/_types.py b/ptodsl/ptodsl/_types.py index 96aea46f38..d44d0f93bb 100644 --- a/ptodsl/ptodsl/_types.py +++ b/ptodsl/ptodsl/_types.py @@ -321,7 +321,17 @@ def _classify_storage_dtype(type_obj): return "compute" if Float8E4M3FNType.isinstance(type_obj) or Float8E5M2Type.isinstance(type_obj): return "storage_only" - if any(_isinstance_pto_type(type_obj, name) for name in ("F8E8M0Type", "HiF8Type", "HiF8x2Type", "F4E1M2x2Type", "F4E2M1x2Type")): + if any( + _isinstance_pto_type(type_obj, name) + for name in ( + "BF16x2Type", + "F8E8M0Type", + "HiF8Type", + "HiF8x2Type", + "F4E1M2x2Type", + "F4E2M1x2Type", + ) + ): return "storage_only" if VectorType.isinstance(type_obj): vec_elem = VectorType(type_obj).element_type @@ -547,6 +557,11 @@ def _int_descriptor(width: int, signedness: str): i16x2 = _DType(lambda: VectorType.get([2], IntegerType.get_signless(16))) i32x2 = _DType(lambda: VectorType.get([2], IntegerType.get_signless(32))) +# ``pto.bf16x2`` is the existing builtin/SIMT vector carrier. VMI uses a +# distinct packed carrier type so that its logical lane count remains the +# number of bf16 pairs rather than the number of scalar bf16 lanes. +_vmi_bf16x2 = _DType(lambda: _pto.BF16x2Type.get()) + # ── Type constructor functions ──────────────────────────────────────────────── diff --git a/ptodsl/ptodsl/_vmi_namespace.py b/ptodsl/ptodsl/_vmi_namespace.py index acbdef830b..f926499af0 100644 --- a/ptodsl/ptodsl/_vmi_namespace.py +++ b/ptodsl/ptodsl/_vmi_namespace.py @@ -31,6 +31,7 @@ VMI_LANE_COUNTS, _ensure_tensor_storage_dtype, _resolve, + _vmi_bf16x2, vmi_mask_type, vmi_vreg_type, ) @@ -170,6 +171,13 @@ def _pointer_element_type(type_obj, *, context: str): def _type_bit_width(type_obj, *, context: str): if IntegerType.isinstance(type_obj): return IntegerType(type_obj).width + if _isinstance_pto_type(type_obj, "BF16x2Type"): + return 32 + if any( + _isinstance_pto_type(type_obj, type_name) + for type_name in ("F4E1M2x2Type", "F4E2M1x2Type") + ): + return 8 if Float8E4M3FNType.isinstance(type_obj) or Float8E5M2Type.isinstance(type_obj): return 8 if F16Type.isinstance(type_obj) or BF16Type.isinstance(type_obj): @@ -183,19 +191,59 @@ def _is_vmi_float_element_type(type_obj) -> bool: return any( cls.isinstance(type_obj) for cls in (BF16Type, F16Type, F32Type, Float8E4M3FNType, Float8E5M2Type) + ) or any( + _isinstance_pto_type(type_obj, type_name) + for type_name in ("BF16x2Type", "F4E1M2x2Type", "F4E2M1x2Type") + ) + + +def _isinstance_pto_type(type_obj, type_name: str) -> bool: + type_cls = getattr(_pto, type_name, None) + if type_cls is None: + return False + try: + return type_cls.isinstance(type_obj) + except Exception: + return False + + +def _is_bf16x2_type(type_obj) -> bool: + return _isinstance_pto_type(type_obj, "BF16x2Type") + + +def _is_f4x2_type(type_obj) -> bool: + return any( + _isinstance_pto_type(type_obj, type_name) + for type_name in ("F4E1M2x2Type", "F4E2M1x2Type") + ) + + +def _validate_vmi_vcvt_bf16x2_pair(source_type, result_type, *, context: str) -> bool: + is_supported_pair = ( + _is_bf16x2_type(source_type) and _is_f4x2_type(result_type) + ) or ( + _is_f4x2_type(source_type) and _is_bf16x2_type(result_type) ) + if is_supported_pair: + return True + if _is_bf16x2_type(source_type) or _is_bf16x2_type(result_type): + raise TypeError( + f"{context} supports bf16x2 only for bf16x2 <-> " + f"f4E1M2x2/f4E2M1x2 conversion; got {source_type} -> {result_type}" + ) + return False -def _normalize_vmi_vcvt_rounding(mode, *, context: str): +def _normalize_vmi_vcvt_rounding(mode, *, context: str, allowed=None): token = mode if not isinstance(token, str): token = str(token) if "." in token: token = token.rsplit(".", 1)[-1] normalized = token.strip().upper() - allowed = {"R", "A", "H", "Z"} - if normalized not in allowed: - expected = ", ".join(sorted(allowed)) + allowed_modes = set(allowed or {"R", "A", "H", "Z"}) + if normalized not in allowed_modes: + expected = ", ".join(sorted(allowed_modes)) raise ValueError( f"{context} does not support rounding {mode!r}; expected one of {expected}" ) @@ -645,6 +693,7 @@ def _emit_reduce( class _VMINamespace: vreg = staticmethod(vmi_vreg_type) mask = staticmethod(vmi_mask_type) + bf16x2 = _vmi_bf16x2 @staticmethod def vload( @@ -991,28 +1040,56 @@ def vcvt( if mask is not None: raise _unsupported_vmi_feature_error("pto.vmi.vcvt", "masked form") result_type = _derive_vcvt_result_type(source, to_dtype, context="pto.vmi.vcvt(...)") + source_type = _as_vmi_vreg_type( + _type_of(source), + context="pto.vmi.vcvt(...)", + ) + is_bf16x2_pair = _validate_vmi_vcvt_bf16x2_pair( + source_type.element_type, + result_type.element_type, + context="pto.vmi.vcvt(...)", + ) + is_bf16x2_to_f4x2 = is_bf16x2_pair and _is_bf16x2_type( + source_type.element_type + ) + is_f4x2_to_bf16x2 = is_bf16x2_pair and _is_f4x2_type( + source_type.element_type + ) if rounding is not None: + if is_f4x2_to_bf16x2: + raise ValueError( + "pto.vmi.vcvt(...) does not support rounding for " + "f4E1M2x2/f4E2M1x2 -> bf16x2 conversion" + ) rounding = _normalize_vmi_vcvt_rounding( rounding, context="pto.vmi.vcvt(..., rounding=...)", + allowed={"R", "A", "F", "Z", "C"} if is_bf16x2_to_f4x2 else None, + ) + elif is_bf16x2_to_f4x2: + rounding = "R" + if is_bf16x2_to_f4x2 and saturate is not None: + raise ValueError( + "pto.vmi.vcvt(...) does not support saturate for bf16x2 -> " + "f4E1M2x2/f4E2M1x2 conversion" + ) + if is_f4x2_to_bf16x2 and saturate is not None: + raise ValueError( + "pto.vmi.vcvt(...) does not support saturate for " + "f4E1M2x2/f4E2M1x2 -> bf16x2 conversion" ) if saturate is None: # The VMI verifier requires explicit "SAT" or "NOSAT" for # narrowing and fp-to-int directions. Default to "SAT" when # the user does not specify. - src_bits = _type_bit_width( - _as_vmi_vreg_type(_type_of(source), context="pto.vmi.vcvt(...)").element_type, - context="pto.vmi.vcvt(...)", - ) + src_bits = _type_bit_width(source_type.element_type, context="pto.vmi.vcvt(...)") dst_bits = _type_bit_width( result_type.element_type, context="pto.vmi.vcvt(...)", ) - src_is_fp = _is_vmi_float_element_type( - _as_vmi_vreg_type(_type_of(source), context="pto.vmi.vcvt(...)").element_type - ) + src_is_fp = _is_vmi_float_element_type(source_type.element_type) dst_is_fp = _is_vmi_float_element_type(result_type.element_type) - if src_bits > dst_bits or (src_is_fp and not dst_is_fp): + if not is_bf16x2_pair and (src_bits > dst_bits or (src_is_fp and not dst_is_fp)): saturate = "SAT" return _call_value( "vcvt", diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 0443de30d7..ab75b1700c 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -2923,6 +2923,116 @@ def vmi_round_r_vcvt_probe(): ) +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_bf16x2_to_f4x2_vcvt_probe(): + src_tile = pto.alloc_tile(shape=[1, 256], dtype=pto.bf16) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + wide = pto.vmi.vload(src_ptr, offset, size=256) + pair = pto.vmi.vinterpret_cast(wide, to_dtype=pto.vmi.bf16x2) + default_e1 = pto.vmi.vcvt(pair, pto.f4e1m2x2) + default_e2 = pto.vmi.vcvt(pair, pto.f4e2m1x2) + round_r = pto.vmi.vcvt(pair, pto.f4e1m2x2, rounding=pto.VcvtRoundMode.R) + round_a = pto.vmi.vcvt(pair, pto.f4e1m2x2, rounding=pto.VcvtRoundMode.A) + round_f = pto.vmi.vcvt(pair, pto.f4e1m2x2, rounding=pto.VcvtRoundMode.F) + round_z = pto.vmi.vcvt(pair, pto.f4e2m1x2, rounding=pto.VcvtRoundMode.Z) + round_c = pto.vmi.vcvt(pair, pto.f4e2m1x2, rounding=pto.VcvtRoundMode.C) + + _ = (default_e1, default_e2, round_r, round_a, round_f, round_z, round_c) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_bf16x2_to_f4x2_vstore_probe(): + src_tile = pto.alloc_tile(shape=[1, 512], dtype=pto.bf16) + dst_e1_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.f4e1m2x2) + dst_e2_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.f4e2m1x2) + src_offset = pto.const(64, dtype=pto.index) + dst_offset = pto.const(16, dtype=pto.index) + + wide = pto.vmi.vload(src_tile.as_ptr(), src_offset, size=128) + pair = pto.vmi.vinterpret_cast(wide, to_dtype=pto.vmi.bf16x2) + out_e1 = pto.vmi.vcvt(pair, pto.f4e1m2x2) + out_e2 = pto.vmi.vcvt(pair, pto.f4e2m1x2) + + pto.vmi.vstore(out_e1, dst_e1_tile.as_ptr(), dst_offset) + pto.vmi.vstore(out_e2, dst_e2_tile.as_ptr(), dst_offset) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_bf16x2_to_f4x2_invalid_rounding_probe(): + src_tile = pto.alloc_tile(shape=[1, 256], dtype=pto.bf16) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + wide = pto.vmi.vload(src_ptr, offset, size=256) + pair = pto.vmi.vinterpret_cast(wide, to_dtype=pto.vmi.bf16x2) + _ = pto.vmi.vcvt(pair, pto.f4e1m2x2, rounding=pto.VcvtRoundMode.H) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_bf16x2_to_f4x2_sat_probe(): + src_tile = pto.alloc_tile(shape=[1, 256], dtype=pto.bf16) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + wide = pto.vmi.vload(src_ptr, offset, size=256) + pair = pto.vmi.vinterpret_cast(wide, to_dtype=pto.vmi.bf16x2) + _ = pto.vmi.vcvt(pair, pto.f4e1m2x2, saturate=pto.VcvtSatMode.SAT) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_bf16x2_to_f4x2_nosat_probe(): + src_tile = pto.alloc_tile(shape=[1, 256], dtype=pto.bf16) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + wide = pto.vmi.vload(src_ptr, offset, size=256) + pair = pto.vmi.vinterpret_cast(wide, to_dtype=pto.vmi.bf16x2) + _ = pto.vmi.vcvt(pair, pto.f4e1m2x2, saturate=pto.VcvtSatMode.NOSAT) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_bf16x2_unsupported_pair_probe(): + src_tile = pto.alloc_tile(shape=[1, 256], dtype=pto.bf16) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + wide = pto.vmi.vload(src_ptr, offset, size=256) + pair = pto.vmi.vinterpret_cast(wide, to_dtype=pto.vmi.bf16x2) + _ = pto.vmi.vcvt(pair, pto.f16) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_f4x2_to_bf16x2_vcvt_probe(): + src_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.f4e1m2x2) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + packed = pto.vmi.vload(src_ptr, offset, size=128) + _ = pto.vmi.vcvt(packed, pto.vmi.bf16x2) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_f4x2_to_bf16x2_rounding_probe(): + src_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.f4e1m2x2) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + packed = pto.vmi.vload(src_ptr, offset, size=128) + _ = pto.vmi.vcvt(packed, pto.vmi.bf16x2, rounding=pto.VcvtRoundMode.R) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_f4x2_to_bf16x2_sat_probe(): + src_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.f4e1m2x2) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + packed = pto.vmi.vload(src_ptr, offset, size=128) + _ = pto.vmi.vcvt(packed, pto.vmi.bf16x2, saturate=pto.VcvtSatMode.SAT) + + @pto.jit(target="a5", backend="vpto", mode="explicit") def vmi_unpack_vload_probe(): src_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.i8) @@ -4115,6 +4225,18 @@ def main() -> None: str(pto.f8e8m0.resolve()) == "!pto.f8E8M0", "pto.f8e8m0 should resolve to the public E8M0 scale type", ) + expect( + str(pto.bf16x2.resolve()) == "vector<2xbf16>", + "pto.bf16x2 should remain the builtin two-lane bf16 vector type", + ) + expect( + str(pto.vmi.bf16x2.resolve()) == "!pto.bf16x2", + "pto.vmi.bf16x2 should resolve to the VMI packed bf16 pair carrier", + ) + expect( + hasattr(pto_types._pto, "BF16x2Type"), + "PTO Python dialect bindings should export BF16x2Type", + ) expect( "hif8" in str(pto.hif8.resolve()), "pto.hif8 should resolve to the public HiF8 type", @@ -6678,6 +6800,119 @@ def _enter_inline_simt_with_resource_attr(): 'rounding = "R"' in vmi_round_r_vcvt_text, "pto.vmi.vcvt should preserve the authored R rounding token for fp32->fp8", ) + vmi_bf16x2_to_f4x2_text = vmi_bf16x2_to_f4x2_vcvt_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + vmi_bf16x2_to_f4x2_text, + "VMI bf16x2-to-f4x2 vcvt specialization", + ) + expect( + "!pto.vmi.vreg<256xbf16>" in vmi_bf16x2_to_f4x2_text, + "VMI bf16x2 vcvt probe should load the source as 256 scalar bf16 lanes", + ) + expect( + "!pto.vmi.vreg<128x!pto.bf16x2>" in vmi_bf16x2_to_f4x2_text, + "VMI vinterpret_cast should form 128 logical bf16x2 lanes", + ) + expect( + "!pto.vmi.vreg<128x!pto.f4E1M2x2>" in vmi_bf16x2_to_f4x2_text + and "!pto.vmi.vreg<128x!pto.f4E2M1x2>" in vmi_bf16x2_to_f4x2_text, + "VMI bf16x2 vcvt should preserve both packed fp4 result element types", + ) + expect( + vmi_bf16x2_to_f4x2_text.count('rounding = "R"') >= 2, + "omitted rounding for bf16x2-to-f4x2 should materialize deterministic R rounding", + ) + for rounding in ("R", "A", "F", "Z", "C"): + expect( + f'rounding = "{rounding}"' in vmi_bf16x2_to_f4x2_text, + f"bf16x2-to-f4x2 should accept and preserve {rounding} rounding", + ) + expect( + 'saturate =' not in vmi_bf16x2_to_f4x2_text, + "bf16x2-to-f4x2 VMI vcvt must not emit a saturate attribute", + ) + vmi_bf16x2_to_f4x2_vstore_text = vmi_bf16x2_to_f4x2_vstore_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + vmi_bf16x2_to_f4x2_vstore_text, + "VMI bf16x2-to-f4x2 vcvt/vstore specialization", + ) + expect( + vmi_bf16x2_to_f4x2_vstore_text.count("pto.vmi.vstore") == 2, + "bf16x2-to-f4x2 conversion results should be storable through VMI vstore", + ) + expect( + "pto.vmi.vstore" in vmi_bf16x2_to_f4x2_vstore_text + and "f4E1M2x2" in vmi_bf16x2_to_f4x2_vstore_text + and "f4E2M1x2" in vmi_bf16x2_to_f4x2_vstore_text, + "VMI vstore should preserve both packed fp4 destination element types", + ) + expect( + "!pto.bf16x2" in vmi_bf16x2_to_f4x2_vstore_text + and "!pto.f4E1M2x2" in vmi_bf16x2_to_f4x2_vstore_text + and "!pto.f4E2M1x2" in vmi_bf16x2_to_f4x2_vstore_text, + "size=128 should preserve the packed bf16x2 and fp4x2 element types", + ) + expect( + vmi_bf16x2_to_f4x2_vstore_text.count('rounding = "R"') == 2, + "bf16x2-to-f4x2 vstore results should use explicit default R rounding", + ) + expect( + "saturate =" not in vmi_bf16x2_to_f4x2_vstore_text, + "bf16x2-to-f4x2 vcvt/vstore must not emit a saturate attribute", + ) + expect_raises( + ValueError, + vmi_bf16x2_to_f4x2_invalid_rounding_probe.compile, + "expected one of A, C, F, R, Z", + ) + for invalid_probe in ( + vmi_bf16x2_to_f4x2_sat_probe, + vmi_bf16x2_to_f4x2_nosat_probe, + ): + expect_raises( + ValueError, + invalid_probe.compile, + "does not support saturate for bf16x2", + ) + for invalid_probe in ( + vmi_bf16x2_unsupported_pair_probe, + ): + expect_raises( + TypeError, + invalid_probe.compile, + "supports bf16x2 only for bf16x2 <->", + ) + vmi_f4x2_to_bf16x2_text = vmi_f4x2_to_bf16x2_vcvt_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + vmi_f4x2_to_bf16x2_text, + "VMI f4x2-to-bf16x2 vcvt specialization", + ) + expect( + "!pto.vmi.vreg<128x!pto.f4E1M2x2>" in vmi_f4x2_to_bf16x2_text, + "VMI f4x2 vcvt probe should load the source as 128 packed f4 lanes", + ) + expect( + "!pto.vmi.vreg<128x!pto.bf16x2>" in vmi_f4x2_to_bf16x2_text, + "VMI f4x2-to-bf16x2 vcvt should widen to 128 logical bf16x2 lanes", + ) + expect( + 'rounding =' not in vmi_f4x2_to_bf16x2_text, + "f4x2-to-bf16x2 VMI vcvt must not emit a rounding attribute", + ) + expect( + "saturate =" not in vmi_f4x2_to_bf16x2_text, + "f4x2-to-bf16x2 VMI vcvt must not emit a saturate attribute", + ) + expect_raises( + ValueError, + vmi_f4x2_to_bf16x2_rounding_probe.compile, + "does not support rounding for", + ) + expect_raises( + ValueError, + vmi_f4x2_to_bf16x2_sat_probe.compile, + "does not support saturate for", + ) unpack_missing_dtype_error = expect_raises( TypeError, vmi_unpack_vload_missing_dtype_probe.compile, diff --git a/python/pto/dialects/pto.py b/python/pto/dialects/pto.py index 161cce5404..eeeb85e057 100644 --- a/python/pto/dialects/pto.py +++ b/python/pto/dialects/pto.py @@ -54,6 +54,7 @@ def _export_optional_cext_symbol(name): HiF8Type = _pto_mod.HiF8Type HiF8x2Type = _pto_mod.HiF8x2Type F8E8M0Type = _pto_mod.F8E8M0Type +BF16x2Type = _pto_mod.BF16x2Type F4E1M2x2Type = _pto_mod.F4E1M2x2Type F4E2M1x2Type = _pto_mod.F4E2M1x2Type TensorViewType = _pto_mod.TensorViewType @@ -234,6 +235,7 @@ def fence_scope_attr_builder(value, context=None): "HiF8Type", "HiF8x2Type", "F8E8M0Type", + "BF16x2Type", "F4E1M2x2Type", "F4E2M1x2Type", "TensorViewType", diff --git a/test/lit/pto/low_precision_type_roundtrip.pto b/test/lit/pto/low_precision_type_roundtrip.pto index a67eab462d..bd91048b6e 100644 --- a/test/lit/pto/low_precision_type_roundtrip.pto +++ b/test/lit/pto/low_precision_type_roundtrip.pto @@ -13,7 +13,8 @@ module { %arg0: !pto.hif8, %arg1: !pto.hif8x2, %arg2: !pto.f4E1M2x2, - %arg3: !pto.f4E2M1x2) { + %arg3: !pto.f4E2M1x2, + %arg4: !pto.bf16x2) { return } } @@ -23,3 +24,4 @@ module { // CHECK-SAME: %arg1: !pto.hif8x2 // CHECK-SAME: %arg2: !pto.f4E1M2x2 // CHECK-SAME: %arg3: !pto.f4E2M1x2 +// CHECK-SAME: %arg4: !pto.bf16x2 diff --git a/test/lit/vmi_new/vmi_bf16x2_compute_invalid.pto b/test/lit/vmi_new/vmi_bf16x2_compute_invalid.pto new file mode 100644 index 0000000000..0053c1f966 --- /dev/null +++ b/test/lit/vmi_new/vmi_bf16x2_compute_invalid.pto @@ -0,0 +1,99 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -split-input-file -verify-diagnostics + +module { + func.func @vadd_bf16x2_invalid( + %lhs: !pto.vmi.vreg<64x!pto.bf16x2>, + %rhs: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'pto.vmi.vadd' op does not support bf16x2 VMI element type; bf16x2 is conversion-only}} + %result = pto.vmi.vadd %lhs, %rhs + : !pto.vmi.vreg<64x!pto.bf16x2>, + !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<64x!pto.bf16x2> + return + } +} + +// ----- + +module { + func.func @vexp_bf16x2_invalid( + %source: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'pto.vmi.vexp' op does not support bf16x2 VMI element type; bf16x2 is conversion-only}} + %result = pto.vmi.vexp %source + : !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<64x!pto.bf16x2> + return + } +} + +// ----- + +module { + func.func @vadds_bf16x2_invalid( + %source: !pto.vmi.vreg<64x!pto.bf16x2>, + %scalar: !pto.bf16x2, + %mask: !pto.vmi.mask<64xpred>) { + // expected-error@+1 {{'pto.vmi.vadds' op does not support bf16x2 VMI element type; bf16x2 is conversion-only}} + %result = pto.vmi.vadds %source, %scalar, %mask + : !pto.vmi.vreg<64x!pto.bf16x2>, !pto.bf16x2, + !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64x!pto.bf16x2> + return + } +} + +// ----- + +module { + func.func @vmula_bf16x2_invalid( + %acc: !pto.vmi.vreg<64x!pto.bf16x2>, + %lhs: !pto.vmi.vreg<64x!pto.bf16x2>, + %rhs: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'pto.vmi.vmula' op does not support bf16x2 VMI element type; bf16x2 is conversion-only}} + %result = pto.vmi.vmula %acc, %lhs, %rhs + : !pto.vmi.vreg<64x!pto.bf16x2>, + !pto.vmi.vreg<64x!pto.bf16x2>, + !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<64x!pto.bf16x2> + return + } +} + +// ----- + +module { + func.func @vcmp_bf16x2_invalid( + %lhs: !pto.vmi.vreg<64x!pto.bf16x2>, + %rhs: !pto.vmi.vreg<64x!pto.bf16x2>, + %seed: !pto.vmi.mask<64xpred>) { + // expected-error@+1 {{'pto.vmi.vcmp' op does not support bf16x2 VMI element type; bf16x2 is conversion-only}} + %result = pto.vmi.vcmp %lhs, %rhs, %seed {cmp = "olt"} + : !pto.vmi.vreg<64x!pto.bf16x2>, + !pto.vmi.vreg<64x!pto.bf16x2>, + !pto.vmi.mask<64xpred> + -> !pto.vmi.mask<64xpred> + return + } +} + +// ----- + +module { + func.func @vcadd_bf16x2_invalid( + %source: !pto.vmi.vreg<64x!pto.bf16x2>, + %mask: !pto.vmi.mask<64xpred>) { + // expected-error@+1 {{'pto.vmi.vcadd' op does not support bf16x2 VMI element type; bf16x2 is conversion-only}} + %result = pto.vmi.vcadd %source, %mask {reassoc} + : !pto.vmi.vreg<64x!pto.bf16x2>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<1x!pto.bf16x2> + return + } +} diff --git a/test/lit/vmi_new/vmi_bf16x2_conversion_pairs_invalid.pto b/test/lit/vmi_new/vmi_bf16x2_conversion_pairs_invalid.pto new file mode 100644 index 0000000000..d5e824a9d0 --- /dev/null +++ b/test/lit/vmi_new/vmi_bf16x2_conversion_pairs_invalid.pto @@ -0,0 +1,129 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -split-input-file -verify-diagnostics + +// bf16x2 is a conversion-only carrier. Phase 1 supports exactly +// bf16x2 <-> f4E1M2x2/f4E2M1x2 (quant narrow and dequant widen); every other +// conversion pair is rejected. + +module { + func.func @unified_bf16x2_to_f16( + %source: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'pto.vmi.vcvt' op unsupported conversion involving bf16x2 element type}} + %result = pto.vmi.vcvt %source + : !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<64xf16> + return + } +} + +// ----- + +module { + func.func @unified_bf16x2_to_bf16( + %source: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'pto.vmi.vcvt' op unsupported conversion involving bf16x2 element type}} + %result = pto.vmi.vcvt %source + : !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<64xbf16> + return + } +} + +// ----- + +module { + func.func @unified_bf16x2_to_f8( + %source: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'pto.vmi.vcvt' op unsupported conversion involving bf16x2 element type}} + %result = pto.vmi.vcvt %source + : !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<64xf8E4M3FN> + return + } +} + +// ----- + +module { + func.func @unified_si32_to_bf16x2( + %source: !pto.vmi.vreg<64xsi32>) { + // expected-error@+1 {{'pto.vmi.vcvt' op unsupported conversion involving bf16x2 element type}} + %result = pto.vmi.vcvt %source + : !pto.vmi.vreg<64xsi32> + -> !pto.vmi.vreg<64x!pto.bf16x2> + return + } +} + +// ----- + +module { + func.func @unified_bf16x2_to_si32( + %source: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'pto.vmi.vcvt' op unsupported conversion involving bf16x2 element type}} + %result = pto.vmi.vcvt %source + : !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<64xsi32> + return + } +} + +// ----- + +module { + func.func @legacy_bf16x2_to_si32( + %source: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'pto.vmi.fptosi' op unsupported fp-to-si conversion element type pair}} + %result = pto.vmi.fptosi %source {saturate = "NOSAT"} + : !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<64xsi32> + return + } +} + +// ----- + +module { + func.func @legacy_bf16x2_to_ui8( + %source: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'pto.vmi.fptoui' op unsupported fp-to-ui conversion element type pair}} + %result = pto.vmi.fptoui %source {saturate = "SAT"} + : !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<64xui8> + return + } +} + +// ----- + +module { + func.func @legacy_truncf_bf16x2_to_f16( + %source: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'pto.vmi.truncf' op unsupported bf16x2 fp-to-fp conversion element type pair}} + %result = pto.vmi.truncf %source {saturate = "SAT"} + : !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<64xf16> + return + } +} + +// ----- + +module { + func.func @legacy_truncf_bf16x2_to_f8( + %source: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'pto.vmi.truncf' op unsupported bf16x2 fp-to-fp conversion element type pair}} + %result = pto.vmi.truncf %source {saturate = "SAT"} + : !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<64xf8E4M3FN> + return + } +} + diff --git a/test/lit/vmi_new/vmi_bitcast_bf16_d4_to_bf16x2_d4_invalid.pto b/test/lit/vmi_new/vmi_bitcast_bf16_d4_to_bf16x2_d4_invalid.pto new file mode 100644 index 0000000000..a7d0a84a88 --- /dev/null +++ b/test/lit/vmi_new/vmi_bitcast_bf16_d4_to_bf16x2_d4_invalid.pto @@ -0,0 +1,27 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not pto-test-opt %s -vmi-lower-unified-to-legacy -pto-validate-vmi-layout-ir 2>&1 | FileCheck %s + +// N2: the bf16 <-> bf16x2 pairing bitcast is a width-changing bitcast whose +// ONLY legal layout is contiguous. A deinterleaved=4 pairing must be rejected +// by the layout-IR validation. This guards the phase-2 blocker: d4 pairing is +// not unlocked until kWidthChangingBitcastLayoutPatterns gains a d(4) row. + +module { + func.func @pair_d4_invalid( + %s: !pto.vmi.vreg<256xbf16, #pto.vmi.layout>) { + %p = pto.vmi.vinterpret_cast %s + : !pto.vmi.vreg<256xbf16, #pto.vmi.layout> + -> !pto.vmi.vreg<128x!pto.bf16x2, #pto.vmi.layout> + return + } +} + +// CHECK: VMI-LAYOUT-CONTRACT: pto.vmi.bitcast has no registered layout support +// CHECK-SAME: width-changing bitcast layout does not match a bitcast layout table row diff --git a/test/lit/vmi_new/vmi_bitcast_bf16x2_total_bits_invalid.pto b/test/lit/vmi_new/vmi_bitcast_bf16x2_total_bits_invalid.pto new file mode 100644 index 0000000000..a080beebf8 --- /dev/null +++ b/test/lit/vmi_new/vmi_bitcast_bf16x2_total_bits_invalid.pto @@ -0,0 +1,20 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @bf16x2_total_bits_mismatch( + %value: !pto.vmi.vreg<64x!pto.bf16x2>) { + %cast = pto.vmi.vinterpret_cast %value + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64xbf16> + return + } +} + +// CHECK: 'pto.vmi.vinterpret_cast' op requires source and result to carry the same total number of bits diff --git a/test/lit/vmi_new/vmi_packed_fp_conversion_pairs_invalid.pto b/test/lit/vmi_new/vmi_packed_fp_conversion_pairs_invalid.pto new file mode 100644 index 0000000000..1bc481fabe --- /dev/null +++ b/test/lit/vmi_new/vmi_packed_fp_conversion_pairs_invalid.pto @@ -0,0 +1,49 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -split-input-file -verify-diagnostics + +// Packed floating-point carrier types must use an explicitly supported +// fp-to-fp conversion contract instead of being admitted by storage width. + +module { + func.func @unified_hif8x2_to_f4( + %source: !pto.vmi.vreg<64x!pto.hif8x2>) { + // expected-error@+1 {{'pto.vmi.vcvt' op unsupported packed fp-to-fp conversion element type pair}} + %result = pto.vmi.vcvt %source + : !pto.vmi.vreg<64x!pto.hif8x2> + -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + return + } +} + +// ----- + +module { + func.func @unified_f4_to_f16( + %source: !pto.vmi.vreg<64x!pto.f4E1M2x2>) { + // expected-error@+1 {{'pto.vmi.vcvt' op unsupported packed fp-to-fp conversion element type pair}} + %result = pto.vmi.vcvt %source + : !pto.vmi.vreg<64x!pto.f4E1M2x2> + -> !pto.vmi.vreg<64xf16> + return + } +} + +// ----- + +module { + func.func @legacy_hif8x2_to_f4( + %source: !pto.vmi.vreg<64x!pto.hif8x2>) { + // expected-error@+1 {{'pto.vmi.truncf' op unsupported packed fp-to-fp conversion element type pair}} + %result = pto.vmi.truncf %source {saturate = "SAT"} + : !pto.vmi.vreg<64x!pto.hif8x2> + -> !pto.vmi.vreg<64x!pto.f4E2M1x2> + return + } +} diff --git a/test/lit/vmi_new/vmi_to_vpto_bitcast.pto b/test/lit/vmi_new/vmi_to_vpto_bitcast.pto index 7532c01e6a..9caf4b3997 100644 --- a/test/lit/vmi_new/vmi_to_vpto_bitcast.pto +++ b/test/lit/vmi_new/vmi_to_vpto_bitcast.pto @@ -15,6 +15,22 @@ module { : !pto.vmi.vreg<128xf32> -> !pto.vmi.vreg<256xi16> return %cast : !pto.vmi.vreg<256xi16> } + + func.func @vmi_to_vpto_bitcast_bf16_to_bf16x2( + %value: !pto.vmi.vreg<128xbf16>) + -> !pto.vmi.vreg<64x!pto.bf16x2> { + %cast = pto.vmi.vinterpret_cast %value + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + return %cast : !pto.vmi.vreg<64x!pto.bf16x2> + } + + func.func @vmi_to_vpto_bitcast_bf16x2_to_bf16( + %value: !pto.vmi.vreg<64x!pto.bf16x2>) + -> !pto.vmi.vreg<128xbf16> { + %cast = pto.vmi.vinterpret_cast %value + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<128xbf16> + return %cast : !pto.vmi.vreg<128xbf16> + } } // CHECK-LABEL: func.func @vmi_to_vpto_bitcast_f32_to_i16( @@ -27,3 +43,19 @@ module { // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast + +// CHECK-LABEL: func.func @vmi_to_vpto_bitcast_bf16_to_bf16x2( +// CHECK-SAME: %[[BF16:[^)]+]]: !pto.vreg<128xbf16> +// CHECK-SAME: -> !pto.vreg<64x!pto.bf16x2> +// CHECK: %[[PAIR:.*]] = pto.vbitcast %[[BF16]] : !pto.vreg<128xbf16> -> !pto.vreg<64x!pto.bf16x2> +// CHECK: return %[[PAIR]] +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. + +// CHECK-LABEL: func.func @vmi_to_vpto_bitcast_bf16x2_to_bf16( +// CHECK-SAME: %[[PAIR_IN:[^)]+]]: !pto.vreg<64x!pto.bf16x2> +// CHECK-SAME: -> !pto.vreg<128xbf16> +// CHECK: %[[BF16_OUT:.*]] = pto.vbitcast %[[PAIR_IN]] : !pto.vreg<64x!pto.bf16x2> -> !pto.vreg<128xbf16> +// CHECK: return %[[BF16_OUT]] +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_extf_f4x2_to_bf16x2_ls4.pto b/test/lit/vmi_new/vmi_to_vpto_extf_f4x2_to_bf16x2_ls4.pto new file mode 100644 index 0000000000..ea2d7c2002 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_extf_f4x2_to_bf16x2_ls4.pto @@ -0,0 +1,38 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s --check-prefix=LOWER + +// Reverse widen (dequant) f4x2 -> bf16x2, compact path: lane-stride-4 source +// (1/4 active f4x2 lanes) -> contiguous bf16x2 result. Single vcvt{P0} 1:1 +// plus the result-side bf16 -> bf16x2 reinterpret. No vor, no rnd, no sat. + +module { + func.func @f4e1m2x2_ls4_to_bf16x2_contiguous( + %src: !pto.vmi.vreg<64x!pto.f4E1M2x2, + #pto.vmi.layout>, + %dst: !pto.ptr, + %off: index) { + %r = pto.vmi.vcvt %src + : !pto.vmi.vreg<64x!pto.f4E1M2x2, + #pto.vmi.layout> + -> !pto.vmi.vreg<64x!pto.bf16x2, #pto.vmi.layout> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<64x!pto.bf16x2, #pto.vmi.layout>, + !pto.ptr + return + } +} + +// LOWER-LABEL: func.func @f4e1m2x2_ls4_to_bf16x2_contiguous( +// LOWER: pto.pset_b8 "PAT_ALL" : !pto.mask +// LOWER: pto.vcvt {{.*}} {part = "P0"} : !pto.vreg<256x!pto.f4E1M2x2>, !pto.mask -> !pto.vreg<128xbf16> +// LOWER: pto.vbitcast {{.*}} : !pto.vreg<128xbf16> -> !pto.vreg<64x!pto.bf16x2> +// LOWER-NOT: pto.vor +// LOWER-NOT: rnd = +// LOWER-NOT: sat = diff --git a/test/lit/vmi_new/vmi_to_vpto_extf_f4x2_to_bf16x2_variants.pto b/test/lit/vmi_new/vmi_to_vpto_extf_f4x2_to_bf16x2_variants.pto new file mode 100644 index 0000000000..3c62c214b1 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_extf_f4x2_to_bf16x2_variants.pto @@ -0,0 +1,70 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s --check-prefix=LOWER + +// Reverse widen (dequant) f4x2 -> bf16x2 variants matrix: 8->32, Packed4, +// no rounding, no saturate. The physical pto.vcvt widens f4 -> bf16; each +// P0-P3 part yields a bf16 register that is reinterpreted to bf16x2 with a +// physical-noop vbitcast (mirror of viewVcvtSource in the truncf direction). +// No vor: the four part results are independent registers. + +module { + // Single source register: 256 f4E1M2x2 (2048bit) -> 256 bf16x2 (8192bit, + // d4): 4 x vcvt{P0-P3} -> 4 x 128 bf16, each vbitcast to 64 bf16x2. + func.func @f4e1m2x2_to_bf16x2( + %src: !pto.vmi.vreg<256x!pto.f4E1M2x2>, + %dst: !pto.ptr, + %off: index) { + %r = pto.vmi.vcvt %src + : !pto.vmi.vreg<256x!pto.f4E1M2x2> -> !pto.vmi.vreg<256x!pto.bf16x2> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<256x!pto.bf16x2>, !pto.ptr + return + } + + // Second packed FP4 format variant. + func.func @f4e2m1x2_to_bf16x2( + %src: !pto.vmi.vreg<256x!pto.f4E2M1x2>, + %dst: !pto.ptr, + %off: index) { + %r = pto.vmi.vcvt %src + : !pto.vmi.vreg<256x!pto.f4E2M1x2> -> !pto.vmi.vreg<256x!pto.bf16x2> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<256x!pto.bf16x2>, !pto.ptr + return + } +} + +// LOWER-LABEL: func.func @f4e1m2x2_to_bf16x2( +// LOWER: pto.pset_b8 "PAT_ALL" : !pto.mask +// LOWER: pto.vcvt {{.*}} {part = "P{{[0-3]}}"} : !pto.vreg<256x!pto.f4E1M2x2>, !pto.mask -> !pto.vreg<128xbf16> +// LOWER: pto.vbitcast {{.*}} : !pto.vreg<128xbf16> -> !pto.vreg<64x!pto.bf16x2> +// LOWER: pto.vcvt {{.*}} {part = "P{{[0-3]}}"} : !pto.vreg<256x!pto.f4E1M2x2>, !pto.mask -> !pto.vreg<128xbf16> +// LOWER: pto.vbitcast {{.*}} : !pto.vreg<128xbf16> -> !pto.vreg<64x!pto.bf16x2> +// LOWER: pto.vcvt {{.*}} {part = "P{{[0-3]}}"} : !pto.vreg<256x!pto.f4E1M2x2>, !pto.mask -> !pto.vreg<128xbf16> +// LOWER: pto.vbitcast {{.*}} : !pto.vreg<128xbf16> -> !pto.vreg<64x!pto.bf16x2> +// LOWER: pto.vcvt {{.*}} {part = "P{{[0-3]}}"} : !pto.vreg<256x!pto.f4E1M2x2>, !pto.mask -> !pto.vreg<128xbf16> +// LOWER: pto.vbitcast {{.*}} : !pto.vreg<128xbf16> -> !pto.vreg<64x!pto.bf16x2> +// LOWER-NOT: pto.vor +// LOWER-NOT: rnd = +// LOWER-NOT: sat = + +// LOWER-LABEL: func.func @f4e2m1x2_to_bf16x2( +// LOWER: pto.pset_b8 "PAT_ALL" : !pto.mask +// LOWER: pto.vcvt {{.*}} {part = "P{{[0-3]}}"} : !pto.vreg<256x!pto.f4E2M1x2>, !pto.mask -> !pto.vreg<128xbf16> +// LOWER: pto.vbitcast {{.*}} : !pto.vreg<128xbf16> -> !pto.vreg<64x!pto.bf16x2> +// LOWER: pto.vcvt {{.*}} {part = "P{{[0-3]}}"} : !pto.vreg<256x!pto.f4E2M1x2>, !pto.mask -> !pto.vreg<128xbf16> +// LOWER: pto.vbitcast {{.*}} : !pto.vreg<128xbf16> -> !pto.vreg<64x!pto.bf16x2> +// LOWER: pto.vcvt {{.*}} {part = "P{{[0-3]}}"} : !pto.vreg<256x!pto.f4E2M1x2>, !pto.mask -> !pto.vreg<128xbf16> +// LOWER: pto.vbitcast {{.*}} : !pto.vreg<128xbf16> -> !pto.vreg<64x!pto.bf16x2> +// LOWER: pto.vcvt {{.*}} {part = "P{{[0-3]}}"} : !pto.vreg<256x!pto.f4E2M1x2>, !pto.mask -> !pto.vreg<128xbf16> +// LOWER: pto.vbitcast {{.*}} : !pto.vreg<128xbf16> -> !pto.vreg<64x!pto.bf16x2> +// LOWER-NOT: pto.vor +// LOWER-NOT: rnd = +// LOWER-NOT: sat = diff --git a/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d2_dynamic_mask.pto b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d2_dynamic_mask.pto new file mode 100644 index 0000000000..23c2220786 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d2_dynamic_mask.pto @@ -0,0 +1,71 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-layout-rematerialize | FileCheck %s --check-prefix=ASSIGN +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-layout-rematerialize -vmi-to-vpto | FileCheck %s --check-prefix=LOWER --implicit-check-not='part = "P1"' --implicit-check-not='part = "P3"' --implicit-check-not=pto.punpack --implicit-check-not=pto.vmi. --implicit-check-not='!pto.vmi.' --implicit-check-not=unrealized_conversion_cast + +// Covers the d2 P0/P2 conversion path together with a runtime tail. Layout +// rematerialization leaves the original contiguous b32 donor predicate +// separate from the lane_stride=2 b16 store predicate. The vcvt and vor +// operations use their own all-true b16 and b8 predicates. + +module { + func.func @bf16x2_d2_dynamic_mask( + %src: !pto.vmi.vreg<128x!pto.bf16x2, + #pto.vmi.layout>, + %dst: !pto.ptr, + %off: index, + %active: index) { + %mask = pto.vmi.create_mask %active + : index -> !pto.vmi.mask<128xpred> + %packed = pto.vmi.vcvt %src {rounding = "Z"} + : !pto.vmi.vreg<128x!pto.bf16x2, + #pto.vmi.layout> + -> !pto.vmi.vreg<128x!pto.f4E1M2x2> + pto.vmi.vstore %packed, %dst[%off], %mask + : !pto.vmi.vreg<128x!pto.f4E1M2x2>, + !pto.ptr, + !pto.vmi.mask<128xpred> + return + } +} + +// ASSIGN-LABEL: func.func @bf16x2_d2_dynamic_mask( +// ASSIGN-SAME: %[[SRC:.*]]: !pto.vmi.vreg<128x!pto.bf16x2, #pto.vmi.layout> +// ASSIGN: %{{.*}} = pto.vmi.create_mask +// ASSIGN-SAME: -> !pto.vmi.mask<128xb32, #pto.vmi.layout> +// ASSIGN: %[[PACKED:.*]] = pto.vmi.truncf %[[SRC]] +// ASSIGN-SAME: -> !pto.vmi.vreg<128x!pto.f4E1M2x2, #pto.vmi.layout> +// ASSIGN: %[[MASK:.*]] = pto.vmi.create_mask +// ASSIGN-SAME: -> !pto.vmi.mask<128xb8, #pto.vmi.layout> +// ASSIGN: pto.vmi.masked_store %[[PACKED]], {{.*}}, %[[MASK]] +// ASSIGN-SAME: !pto.vmi.mask<128xb8, #pto.vmi.layout> + +// LOWER-LABEL: func.func @bf16x2_d2_dynamic_mask( +// LOWER-SAME: %[[P0:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// LOWER-SAME: %[[P2:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// LOWER-SAME: %[[DST:[^,]+]]: !pto.ptr +// LOWER-SAME: %[[OFF:[^,]+]]: index +// LOWER-SAME: %[[ACTIVE:[^)]+]]: index) +// LOWER: %[[ACTIVE_I32:.*]] = arith.index_cast %[[ACTIVE]] : index to i32 +// LOWER: %[[NONNEG:.*]] = arith.maxsi %[[ACTIVE_I32]], {{.*}} : i32 +// LOWER: %[[CLAMPED:.*]] = arith.minui %[[NONNEG]], {{.*}} : i32 +// LOWER: %[[SRC_M0:.*]], %[[SRC_REM0:.*]] = pto.plt_b32 %[[CLAMPED]] : i32 -> !pto.mask, i32 +// LOWER: %[[SRC_M1:.*]], %{{.*}} = pto.plt_b32 %[[SRC_REM0]] : i32 -> !pto.mask, i32 +// LOWER-DAG: %[[VCVT_MASK:.*]] = pto.pset_b16 "PAT_ALL" : !pto.mask +// LOWER-DAG: %[[MERGE_MASK:.*]] = pto.pset_b8 "PAT_ALL" : !pto.mask +// LOWER: %[[P0_BF16:.*]] = pto.vbitcast %[[P0]] : !pto.vreg<64x!pto.bf16x2> -> !pto.vreg<128xbf16> +// LOWER: %[[R0:.*]] = pto.vcvt %[[P0_BF16]], %[[VCVT_MASK]] {part = "P0", rnd = "Z"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: %[[P2_BF16:.*]] = pto.vbitcast %[[P2]] : !pto.vreg<64x!pto.bf16x2> -> !pto.vreg<128xbf16> +// LOWER: %[[R2:.*]] = pto.vcvt %[[P2_BF16]], %[[VCVT_MASK]] {part = "P2", rnd = "Z"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: %[[PACKED:.*]] = pto.vor %[[R0]], %[[R2]], %[[MERGE_MASK]] : !pto.vreg<256x!pto.f4E1M2x2>, !pto.vreg<256x!pto.f4E1M2x2>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: %[[STORE_ACTIVE_I32:.*]] = arith.index_cast %[[ACTIVE]] : index to i32 +// LOWER: %[[STORE_NONNEG:.*]] = arith.maxsi %[[STORE_ACTIVE_I32]], {{.*}} : i32 +// LOWER: %[[STORE_CLAMPED:.*]] = arith.minui %[[STORE_NONNEG]], {{.*}} : i32 +// LOWER: %[[STORE_M:.*]], %{{.*}} = pto.plt_b16 %[[STORE_CLAMPED]] : i32 -> !pto.mask, i32 +// LOWER: pto.vsts %[[PACKED]], %[[DST]][%[[OFF]]], %[[STORE_M]] {dist = "PK_B16"} : !pto.vreg<256x!pto.f4E1M2x2>, !pto.ptr, !pto.mask diff --git a/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d2_lane_stride2.pto b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d2_lane_stride2.pto new file mode 100644 index 0000000000..a7951ef5d7 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d2_lane_stride2.pto @@ -0,0 +1,42 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment | FileCheck %s --check-prefix=ASSIGN +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s --check-prefix=LOWER --implicit-check-not='part = "P1"' --implicit-check-not='part = "P3"' --implicit-check-not=pto.vmi. --implicit-check-not='!pto.vmi.' --implicit-check-not=unrealized_conversion_cast + +// The legal 32-to-8 layout row d2 -> ls2 uses the P0 and P2 packed quarters. + +module { + func.func @bf16x2_d2_to_f4x2_lane_stride2( + %src: !pto.vmi.vreg<128x!pto.bf16x2, #pto.vmi.layout>, + %dst: !pto.ptr, + %off: index) { + %r = pto.vmi.vcvt %src {rounding = "A"} + : !pto.vmi.vreg<128x!pto.bf16x2, #pto.vmi.layout> + -> !pto.vmi.vreg<128x!pto.f4E1M2x2> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<128x!pto.f4E1M2x2>, !pto.ptr + return + } +} + +// ASSIGN-LABEL: func.func @bf16x2_d2_to_f4x2_lane_stride2( +// ASSIGN: pto.vmi.truncf +// ASSIGN-SAME: !pto.vmi.vreg<128x!pto.bf16x2, #pto.vmi.layout> -> !pto.vmi.vreg<128x!pto.f4E1M2x2, #pto.vmi.layout> + +// LOWER-LABEL: func.func @bf16x2_d2_to_f4x2_lane_stride2( +// LOWER-SAME: %[[P0:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// LOWER-SAME: %[[P2:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// LOWER-SAME: %[[DST:[^,]+]]: !pto.ptr +// LOWER-SAME: %[[OFF:[^)]+]]: index) +// LOWER: %[[P0_BF16:.*]] = pto.vbitcast %[[P0]] : !pto.vreg<64x!pto.bf16x2> -> !pto.vreg<128xbf16> +// LOWER: %[[R0:.*]] = pto.vcvt %[[P0_BF16]], {{.*}} {part = "P0", rnd = "A"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: %[[P2_BF16:.*]] = pto.vbitcast %[[P2]] : !pto.vreg<64x!pto.bf16x2> -> !pto.vreg<128xbf16> +// LOWER: %[[R2:.*]] = pto.vcvt %[[P2_BF16]], {{.*}} {part = "P2", rnd = "A"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: %[[PACKED:.*]] = pto.vor %[[R0]], %[[R2]] +// LOWER: pto.vsts %[[PACKED]], %[[DST]][%[[OFF]]], {{.*}} {dist = "PK_B16"} : !pto.vreg<256x!pto.f4E1M2x2>, !pto.ptr, !pto.mask diff --git a/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d2_multichunk.pto b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d2_multichunk.pto new file mode 100644 index 0000000000..be6e0e8273 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d2_multichunk.pto @@ -0,0 +1,57 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment | FileCheck %s --check-prefix=ASSIGN +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s --check-prefix=LOWER --implicit-check-not='part = "P1"' --implicit-check-not='part = "P3"' --implicit-check-not=pto.vmi. --implicit-check-not='!pto.vmi.' --implicit-check-not=unrealized_conversion_cast + +// Protect part-major source indexing for a two-chunk d2 conversion. Output +// chunk 0 consumes P0C0/P2C0, output chunk 1 consumes P0C1/P2C1, and the +// compact stores advance by 128 logical f4x2 lanes. + +module { + func.func @bf16x2_d2_multichunk( + %src: !pto.vmi.vreg<256x!pto.bf16x2, + #pto.vmi.layout>, + %dst: !pto.ptr, + %off: index) { + %packed = pto.vmi.vcvt %src {rounding = "C"} + : !pto.vmi.vreg<256x!pto.bf16x2, + #pto.vmi.layout> + -> !pto.vmi.vreg<256x!pto.f4E2M1x2> + pto.vmi.vstore %packed, %dst[%off] + : !pto.vmi.vreg<256x!pto.f4E2M1x2>, + !pto.ptr + return + } +} + +// ASSIGN-LABEL: func.func @bf16x2_d2_multichunk( +// ASSIGN: pto.vmi.truncf +// ASSIGN-SAME: !pto.vmi.vreg<256x!pto.bf16x2, #pto.vmi.layout> -> !pto.vmi.vreg<256x!pto.f4E2M1x2, #pto.vmi.layout> + +// LOWER-LABEL: func.func @bf16x2_d2_multichunk( +// LOWER-SAME: %[[P0C0:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// LOWER-SAME: %[[P0C1:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// LOWER-SAME: %[[P2C0:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// LOWER-SAME: %[[P2C1:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// LOWER-SAME: %[[DST:[^,]+]]: !pto.ptr +// LOWER-SAME: %[[OFF:[^)]+]]: index) +// LOWER: %[[P0C0_BF16:.*]] = pto.vbitcast %[[P0C0]] +// LOWER: %[[C0P0:.*]] = pto.vcvt %[[P0C0_BF16]], {{.*}} {part = "P0", rnd = "C"} +// LOWER: %[[P2C0_BF16:.*]] = pto.vbitcast %[[P2C0]] +// LOWER: %[[C0P2:.*]] = pto.vcvt %[[P2C0_BF16]], {{.*}} {part = "P2", rnd = "C"} +// LOWER: %[[C0:.*]] = pto.vor %[[C0P0]], %[[C0P2]] +// LOWER: %[[P0C1_BF16:.*]] = pto.vbitcast %[[P0C1]] +// LOWER: %[[C1P0:.*]] = pto.vcvt %[[P0C1_BF16]], {{.*}} {part = "P0", rnd = "C"} +// LOWER: %[[P2C1_BF16:.*]] = pto.vbitcast %[[P2C1]] +// LOWER: %[[C1P2:.*]] = pto.vcvt %[[P2C1_BF16]], {{.*}} {part = "P2", rnd = "C"} +// LOWER: %[[C1:.*]] = pto.vor %[[C1P0]], %[[C1P2]] +// LOWER: pto.vsts %[[C0]], %[[DST]][%[[OFF]]], {{.*}} {dist = "PK_B16"} +// LOWER: %[[C128:.*]] = arith.constant 128 : index +// LOWER: %[[OFF128:.*]] = arith.addi %[[OFF]], %[[C128]] : index +// LOWER: pto.vsts %[[C1]], %[[DST]][%[[OFF128]]], {{.*}} {dist = "PK_B16"} diff --git a/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d4_dynamic_mask.pto b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d4_dynamic_mask.pto new file mode 100644 index 0000000000..c50aa1aaa6 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d4_dynamic_mask.pto @@ -0,0 +1,81 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-layout-rematerialize | FileCheck %s --check-prefix=ASSIGN +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-layout-rematerialize -vmi-to-vpto | FileCheck %s --check-prefix=LOWER --implicit-check-not='dist = "PK4_B32"' --implicit-check-not=pto.vmi. --implicit-check-not='!pto.vmi.' --implicit-check-not=unrealized_conversion_cast + +// Crosses the d4 Packed4 conversion path with a runtime tail and a compact +// contiguous store. The source is pre-annotated d4 so this isolates conversion, +// the original contiguous donor mask, and store-mask rematerialization from +// the d4 pairing-bitcast restriction. + +module { + func.func @bf16x2_d4_dynamic_mask( + %src: !pto.vmi.vreg<256x!pto.bf16x2, + #pto.vmi.layout>, + %dst: !pto.ptr, + %off: index, + %active: index) { + %mask = pto.vmi.create_mask %active + : index -> !pto.vmi.mask<256xpred> + %packed = pto.vmi.vcvt %src {rounding = "R"} + : !pto.vmi.vreg<256x!pto.bf16x2, + #pto.vmi.layout> + -> !pto.vmi.vreg<256x!pto.f4E1M2x2> + pto.vmi.vstore %packed, %dst[%off], %mask + : !pto.vmi.vreg<256x!pto.f4E1M2x2>, + !pto.ptr, + !pto.vmi.mask<256xpred> + return + } +} + +// ASSIGN-LABEL: func.func @bf16x2_d4_dynamic_mask( +// ASSIGN-SAME: %[[SRC:.*]]: !pto.vmi.vreg<256x!pto.bf16x2, #pto.vmi.layout> +// ASSIGN: %{{.*}} = pto.vmi.create_mask +// ASSIGN-SAME: -> !pto.vmi.mask<256xb32, #pto.vmi.layout> +// ASSIGN: %[[PACKED:.*]] = pto.vmi.truncf %[[SRC]] +// ASSIGN-SAME: -> !pto.vmi.vreg<256x!pto.f4E1M2x2, #pto.vmi.layout> +// ASSIGN: %[[MASK:.*]] = pto.vmi.create_mask +// ASSIGN-SAME: -> !pto.vmi.mask<256xb8, #pto.vmi.layout> +// ASSIGN: pto.vmi.masked_store %[[PACKED]], {{.*}}, %[[MASK]] +// ASSIGN-SAME: !pto.vmi.mask<256xb8, #pto.vmi.layout> + +// LOWER-LABEL: func.func @bf16x2_d4_dynamic_mask( +// LOWER-SAME: %[[P0:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// LOWER-SAME: %[[P1:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// LOWER-SAME: %[[P2:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// LOWER-SAME: %[[P3:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// LOWER-SAME: %[[DST:[^,]+]]: !pto.ptr +// LOWER-SAME: %[[OFF:[^,]+]]: index +// LOWER-SAME: %[[ACTIVE:[^)]+]]: index) +// LOWER: %[[ACTIVE_I32:.*]] = arith.index_cast %[[ACTIVE]] : index to i32 +// LOWER: %[[NONNEG:.*]] = arith.maxsi %[[ACTIVE_I32]], {{.*}} : i32 +// LOWER: %[[CLAMPED:.*]] = arith.minui %[[NONNEG]], {{.*}} : i32 +// LOWER: %[[SRC_M0:.*]], %[[SRC_REM0:.*]] = pto.plt_b32 %[[CLAMPED]] : i32 -> !pto.mask, i32 +// LOWER: %[[SRC_M1:.*]], %[[SRC_REM1:.*]] = pto.plt_b32 %[[SRC_REM0]] : i32 -> !pto.mask, i32 +// LOWER: %[[SRC_M2:.*]], %[[SRC_REM2:.*]] = pto.plt_b32 %[[SRC_REM1]] : i32 -> !pto.mask, i32 +// LOWER: %[[SRC_M3:.*]], %{{.*}} = pto.plt_b32 %[[SRC_REM2]] : i32 -> !pto.mask, i32 +// LOWER-DAG: %[[VCVT_MASK:.*]] = pto.pset_b16 "PAT_ALL" : !pto.mask +// LOWER-DAG: %[[MERGE_MASK:.*]] = pto.pset_b8 "PAT_ALL" : !pto.mask +// LOWER: %[[P0_BF16:.*]] = pto.vbitcast %[[P0]] : !pto.vreg<64x!pto.bf16x2> -> !pto.vreg<128xbf16> +// LOWER: %[[R0:.*]] = pto.vcvt %[[P0_BF16]], %[[VCVT_MASK]] {part = "P0", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: %[[P1_BF16:.*]] = pto.vbitcast %[[P1]] : !pto.vreg<64x!pto.bf16x2> -> !pto.vreg<128xbf16> +// LOWER: %[[R1:.*]] = pto.vcvt %[[P1_BF16]], %[[VCVT_MASK]] {part = "P1", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: %[[P2_BF16:.*]] = pto.vbitcast %[[P2]] : !pto.vreg<64x!pto.bf16x2> -> !pto.vreg<128xbf16> +// LOWER: %[[R2:.*]] = pto.vcvt %[[P2_BF16]], %[[VCVT_MASK]] {part = "P2", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: %[[P3_BF16:.*]] = pto.vbitcast %[[P3]] : !pto.vreg<64x!pto.bf16x2> -> !pto.vreg<128xbf16> +// LOWER: %[[R3:.*]] = pto.vcvt %[[P3_BF16]], %[[VCVT_MASK]] {part = "P3", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: %[[M1:.*]] = pto.vor %[[R0]], %[[R1]], %[[MERGE_MASK]] +// LOWER: %[[M2:.*]] = pto.vor %[[M1]], %[[R2]], %[[MERGE_MASK]] +// LOWER: %[[PACKED:.*]] = pto.vor %[[M2]], %[[R3]], %[[MERGE_MASK]] +// LOWER: %[[STORE_ACTIVE_I32:.*]] = arith.index_cast %[[ACTIVE]] : index to i32 +// LOWER: %[[STORE_NONNEG:.*]] = arith.maxsi %[[STORE_ACTIVE_I32]], {{.*}} : i32 +// LOWER: %[[STORE_CLAMPED:.*]] = arith.minui %[[STORE_NONNEG]], {{.*}} : i32 +// LOWER: %[[STORE_M:.*]], %{{.*}} = pto.plt_b8 %[[STORE_CLAMPED]] : i32 -> !pto.mask, i32 +// LOWER: pto.vsts %[[PACKED]], %[[DST]][%[[OFF]]], %[[STORE_M]] : !pto.vreg<256x!pto.f4E1M2x2>, !pto.ptr, !pto.mask diff --git a/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d4_multichunk.pto b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d4_multichunk.pto new file mode 100644 index 0000000000..b5ca963ab7 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d4_multichunk.pto @@ -0,0 +1,59 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s --implicit-check-not=pto.vmi. --implicit-check-not='!pto.vmi.' --implicit-check-not=unrealized_conversion_cast + +// Protect the part-major physical source order for a two-chunk d4 conversion. +// Each output chunk must consume the matching chunk from P0, P1, P2, and P3. + +module { + func.func @bf16x2_d4_multichunk( + %src: !pto.vmi.vreg<512x!pto.bf16x2, + #pto.vmi.layout>) + -> !pto.vmi.vreg<512x!pto.f4E2M1x2> { + %r = pto.vmi.vcvt %src {rounding = "R"} + : !pto.vmi.vreg<512x!pto.bf16x2, + #pto.vmi.layout> + -> !pto.vmi.vreg<512x!pto.f4E2M1x2> + return %r : !pto.vmi.vreg<512x!pto.f4E2M1x2> + } +} + +// CHECK-LABEL: func.func @bf16x2_d4_multichunk( +// CHECK-SAME: %[[P0C0:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// CHECK-SAME: %[[P0C1:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// CHECK-SAME: %[[P1C0:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// CHECK-SAME: %[[P1C1:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// CHECK-SAME: %[[P2C0:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// CHECK-SAME: %[[P2C1:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// CHECK-SAME: %[[P3C0:[^,]+]]: !pto.vreg<64x!pto.bf16x2> +// CHECK-SAME: %[[P3C1:[^)]+]]: !pto.vreg<64x!pto.bf16x2> +// CHECK-SAME: -> (!pto.vreg<256x!pto.f4E2M1x2>, !pto.vreg<256x!pto.f4E2M1x2>) +// CHECK: %[[P0C0_BF16:.*]] = pto.vbitcast %[[P0C0]] +// CHECK: %[[C0P0:.*]] = pto.vcvt %[[P0C0_BF16]], {{.*}} {part = "P0", rnd = "R"} +// CHECK: %[[P1C0_BF16:.*]] = pto.vbitcast %[[P1C0]] +// CHECK: %[[C0P1:.*]] = pto.vcvt %[[P1C0_BF16]], {{.*}} {part = "P1", rnd = "R"} +// CHECK: %[[P2C0_BF16:.*]] = pto.vbitcast %[[P2C0]] +// CHECK: %[[C0P2:.*]] = pto.vcvt %[[P2C0_BF16]], {{.*}} {part = "P2", rnd = "R"} +// CHECK: %[[P3C0_BF16:.*]] = pto.vbitcast %[[P3C0]] +// CHECK: %[[C0P3:.*]] = pto.vcvt %[[P3C0_BF16]], {{.*}} {part = "P3", rnd = "R"} +// CHECK: %[[C0M1:.*]] = pto.vor %[[C0P0]], %[[C0P1]] +// CHECK: %[[C0M2:.*]] = pto.vor %[[C0M1]], %[[C0P2]] +// CHECK: %[[C0:.*]] = pto.vor %[[C0M2]], %[[C0P3]] +// CHECK: %[[P0C1_BF16:.*]] = pto.vbitcast %[[P0C1]] +// CHECK: %[[C1P0:.*]] = pto.vcvt %[[P0C1_BF16]], {{.*}} {part = "P0", rnd = "R"} +// CHECK: %[[P1C1_BF16:.*]] = pto.vbitcast %[[P1C1]] +// CHECK: %[[C1P1:.*]] = pto.vcvt %[[P1C1_BF16]], {{.*}} {part = "P1", rnd = "R"} +// CHECK: %[[P2C1_BF16:.*]] = pto.vbitcast %[[P2C1]] +// CHECK: %[[C1P2:.*]] = pto.vcvt %[[P2C1_BF16]], {{.*}} {part = "P2", rnd = "R"} +// CHECK: %[[P3C1_BF16:.*]] = pto.vbitcast %[[P3C1]] +// CHECK: %[[C1P3:.*]] = pto.vcvt %[[P3C1_BF16]], {{.*}} {part = "P3", rnd = "R"} +// CHECK: %[[C1M1:.*]] = pto.vor %[[C1P0]], %[[C1P1]] +// CHECK: %[[C1M2:.*]] = pto.vor %[[C1M1]], %[[C1P2]] +// CHECK: %[[C1:.*]] = pto.vor %[[C1M2]], %[[C1P3]] +// CHECK: return %[[C0]], %[[C1]] diff --git a/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d4_packed4.pto b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d4_packed4.pto new file mode 100644 index 0000000000..aea18ef2f9 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_d4_packed4.pto @@ -0,0 +1,38 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s --check-prefix=LOWER + +// L3: d4 direct truncf probe. A pre-annotated deinterleaved=4 bf16x2 source +// feeds the full Packed4 path (4x vcvt{P0-P3} + 3x vor). NOTE: this uses a +// pre-annotated argument to bypass the pairing bitcast (which is only legal +// under contiguous; see the d4 pairing negative test). It is a lowering +// probe, not the full real-data flow (which needs phase-2 pairing unlock). + +module { + func.func @bf16x2_d4_to_f4x2_packed4( + %src: !pto.vmi.vreg<256x!pto.bf16x2, #pto.vmi.layout>, + %dst: !pto.ptr, + %off: index) { + %r = pto.vmi.vcvt %src {rounding = "R"} + : !pto.vmi.vreg<256x!pto.bf16x2, #pto.vmi.layout> + -> !pto.vmi.vreg<256x!pto.f4E1M2x2> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<256x!pto.f4E1M2x2>, !pto.ptr + return + } +} + +// LOWER-LABEL: func.func @bf16x2_d4_to_f4x2_packed4( +// LOWER: pto.vcvt {{.*}} {part = "P0", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: pto.vcvt {{.*}} {part = "P1", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: pto.vcvt {{.*}} {part = "P2", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: pto.vcvt {{.*}} {part = "P3", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER-COUNT-3: pto.vor +// LOWER-NOT: pto.vmi. +// LOWER-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_dynamic_tail.pto b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_dynamic_tail.pto new file mode 100644 index 0000000000..41bd7674a4 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_dynamic_tail.pto @@ -0,0 +1,73 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-layout-rematerialize | FileCheck %s --check-prefix=ASSIGN +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-layout-rematerialize -vmi-to-vpto | FileCheck %s --check-prefix=LOWER --implicit-check-not='part = "P1"' --implicit-check-not='part = "P2"' --implicit-check-not='part = "P3"' --implicit-check-not=pto.vor --implicit-check-not=pto.vmi. --implicit-check-not='!pto.vmi.' --implicit-check-not=unrealized_conversion_cast + +// Combination guard for a runtime tail count. The original contiguous b32 mask +// and the lane_stride=4 store mask are materialized independently. Conversion +// uses an all-true b16 predicate, followed by two masked PK4_B32 stores without +// compact-result assembly. + +module { + func.func @bf16x2_dynamic_tail( + %src: !pto.vmi.vreg<256xbf16>, + %dst: !pto.ptr, + %off: index, + %active: index) { + %mask = pto.vmi.create_mask %active + : index -> !pto.vmi.mask<128xpred> + %pair = pto.vmi.vinterpret_cast %src + : !pto.vmi.vreg<256xbf16> -> !pto.vmi.vreg<128x!pto.bf16x2> + %packed = pto.vmi.vcvt %pair {rounding = "R"} + : !pto.vmi.vreg<128x!pto.bf16x2> + -> !pto.vmi.vreg<128x!pto.f4E1M2x2> + pto.vmi.vstore %packed, %dst[%off], %mask + : !pto.vmi.vreg<128x!pto.f4E1M2x2>, + !pto.ptr, + !pto.vmi.mask<128xpred> + return + } +} + +// ASSIGN-LABEL: func.func @bf16x2_dynamic_tail( +// ASSIGN-SAME: %[[SRC:.*]]: !pto.vmi.vreg<256xbf16, #pto.vmi.layout> +// ASSIGN: %{{.*}} = pto.vmi.create_mask +// ASSIGN-SAME: -> !pto.vmi.mask<128xb32, #pto.vmi.layout> +// ASSIGN: %[[PAIR:.*]] = pto.vmi.bitcast %[[SRC]] +// ASSIGN-SAME: -> !pto.vmi.vreg<128x!pto.bf16x2, #pto.vmi.layout> +// ASSIGN: %[[PACKED:.*]] = pto.vmi.truncf %[[PAIR]] +// ASSIGN-SAME: -> !pto.vmi.vreg<128x!pto.f4E1M2x2, #pto.vmi.layout> +// ASSIGN: %[[MASK:.*]] = pto.vmi.create_mask +// ASSIGN-SAME: -> !pto.vmi.mask<128xb8, #pto.vmi.layout> +// ASSIGN: pto.vmi.masked_store %[[PACKED]], {{.*}}, %[[MASK]] +// ASSIGN-SAME: !pto.vmi.mask<128xb8, #pto.vmi.layout> + +// LOWER-LABEL: func.func @bf16x2_dynamic_tail( +// LOWER-SAME: %[[SRC0:[^,]+]]: !pto.vreg<128xbf16> +// LOWER-SAME: %[[SRC1:[^,]+]]: !pto.vreg<128xbf16> +// LOWER-SAME: %[[DST:[^,]+]]: !pto.ptr +// LOWER-SAME: %[[OFF:[^,]+]]: index +// LOWER-SAME: %[[ACTIVE:[^)]+]]: index) +// LOWER: %[[ACTIVE_I32:.*]] = arith.index_cast %[[ACTIVE]] : index to i32 +// LOWER: %[[NONNEG:.*]] = arith.maxsi %[[ACTIVE_I32]], {{.*}} : i32 +// LOWER: %[[CLAMPED:.*]] = arith.minui %[[NONNEG]], {{.*}} : i32 +// LOWER: %[[SRC_M0:.*]], %[[SRC_REM:.*]] = pto.plt_b32 %[[CLAMPED]] : i32 -> !pto.mask, i32 +// LOWER: %[[SRC_M1:.*]], %{{.*}} = pto.plt_b32 %[[SRC_REM]] : i32 -> !pto.mask, i32 +// LOWER-DAG: %[[VCVT_MASK:.*]] = pto.pset_b16 "PAT_ALL" : !pto.mask +// LOWER: %[[R0:.*]] = pto.vcvt %[[SRC0]], %[[VCVT_MASK]] {part = "P0", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: %[[R1:.*]] = pto.vcvt %[[SRC1]], %[[VCVT_MASK]] {part = "P0", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER: %[[STORE_ACTIVE_I32:.*]] = arith.index_cast %[[ACTIVE]] : index to i32 +// LOWER: %[[STORE_NONNEG:.*]] = arith.maxsi %[[STORE_ACTIVE_I32]], {{.*}} : i32 +// LOWER: %[[STORE_CLAMPED:.*]] = arith.minui %[[STORE_NONNEG]], {{.*}} : i32 +// LOWER: %[[STORE_M0:.*]], %[[STORE_REM:.*]] = pto.plt_b32 %[[STORE_CLAMPED]] : i32 -> !pto.mask, i32 +// LOWER: %[[STORE_M1:.*]], %{{.*}} = pto.plt_b32 %[[STORE_REM]] : i32 -> !pto.mask, i32 +// LOWER: pto.vsts %[[R0]], %[[DST]][%[[OFF]]], %[[STORE_M0]] {dist = "PK4_B32"} : !pto.vreg<256x!pto.f4E1M2x2>, !pto.ptr, !pto.mask +// LOWER: %[[C64:.*]] = arith.constant 64 : index +// LOWER: %[[OFF64:.*]] = arith.addi %[[OFF]], %[[C64]] : index +// LOWER: pto.vsts %[[R1]], %[[DST]][%[[OFF64]]], %[[STORE_M1]] {dist = "PK4_B32"} : !pto.vreg<256x!pto.f4E1M2x2>, !pto.ptr, !pto.mask diff --git a/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_mask_granularity.pto b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_mask_granularity.pto new file mode 100644 index 0000000000..f98be2f283 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_mask_granularity.pto @@ -0,0 +1,35 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s --check-prefix=LOWER + +// L2: the bf16x2 physical source is consumed by pto.vcvt as raw bf16 lanes, +// so the source mask granularity must be b16 (not b32, the bf16x2 storage +// granularity). A regression on this specific point (mask built from the +// bf16x2 view instead of the bf16 view) would silently double the mask width. + +module { + func.func @bf16x2_mask_granularity( + %src: !pto.vmi.vreg<128xbf16>, + %dst: !pto.ptr, + %off: index) { + %pair = pto.vmi.vinterpret_cast %src + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + %r = pto.vmi.vcvt %pair {rounding = "R"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<64x!pto.f4E1M2x2>, !pto.ptr + return + } +} + +// LOWER-LABEL: func.func @bf16x2_mask_granularity( +// LOWER: pto.vcvt {{.*}} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER-NOT: pto.vor +// (The downstream vsts legitimately uses a b32 mask for the packed store, so +// the b16-granularity guard is the positive vcvt check above.) diff --git a/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_to_f4x2_contiguous.pto b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_to_f4x2_contiguous.pto new file mode 100644 index 0000000000..40b0beef58 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_to_f4x2_contiguous.pto @@ -0,0 +1,44 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment | FileCheck %s --check-prefix=ASSIGN +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s --check-prefix=LOWER + +// Phase-1 bf16x2 -> f4x2 narrowing, contiguous path: +// vinterpret_cast 128xbf16 -> 64xbf16x2 (physical noop pairing), +// vcvt 64xbf16x2 -> 64xf4E1M2x2 (32->8, c()->ls(4), single vcvt{P0}, no sat). +// The physical pto.vcvt consumes the bf16x2 source viewed as raw bf16 lanes. + +module { + func.func @bf16x2_to_f4x2_contiguous( + %src: !pto.vmi.vreg<128xbf16>, + %dst: !pto.ptr, + %off: index) { + %pair = pto.vmi.vinterpret_cast %src + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + %r = pto.vmi.vcvt %pair {rounding = "R"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<64x!pto.f4E1M2x2>, !pto.ptr + return + } +} + +// ASSIGN-LABEL: func.func @bf16x2_to_f4x2_contiguous( +// ASSIGN-SAME: %[[INPUT:.*]]: !pto.vmi.vreg<128xbf16, #pto.vmi.layout> +// ASSIGN: %[[PAIR:.*]] = pto.vmi.bitcast %[[INPUT]] +// ASSIGN-SAME: !pto.vmi.vreg<128xbf16, #pto.vmi.layout> -> !pto.vmi.vreg<64x!pto.bf16x2, #pto.vmi.layout> +// ASSIGN: %[[R:.*]] = pto.vmi.truncf %[[PAIR]] +// ASSIGN-SAME: !pto.vmi.vreg<64x!pto.bf16x2, #pto.vmi.layout> -> !pto.vmi.vreg<64x!pto.f4E1M2x2, #pto.vmi.layout> + +// LOWER-LABEL: func.func @bf16x2_to_f4x2_contiguous( +// LOWER: pto.vcvt {{.*}} {part = "P0", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER-NOT: pto.vor +// LOWER-NOT: pto.vmi. +// LOWER-NOT: !pto.vmi. +// LOWER-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_to_f4x2_variants.pto b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_to_f4x2_variants.pto new file mode 100644 index 0000000000..02257adf69 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_truncf_bf16x2_to_f4x2_variants.pto @@ -0,0 +1,123 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s --check-prefix=LOWER + +// L1: contiguous bf16x2 -> f4x2 variants matrix. +// The conversion contract accepts {R,A,F,Z,C}; this matrix exercises every +// accepted token plus the omitted-attribute default (R). + +module { + // Single register: 128 bf16 -> 64 bf16x2 -> 64 f4E1M2x2 (1 physical reg, 1 vcvt{P0}). + func.func @bf16x2_to_f4e1m2x2_r( + %src: !pto.vmi.vreg<128xbf16>, + %dst: !pto.ptr, + %off: index) { + %pair = pto.vmi.vinterpret_cast %src + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + %r = pto.vmi.vcvt %pair {rounding = "R"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<64x!pto.f4E1M2x2>, !pto.ptr + return + } + + // Multi register: 256 bf16 -> 128 bf16x2 -> 128 f4E2M1x2 (2 physical regs, 2 vcvt{P0}). + func.func @bf16x2_to_f4e2m1x2_a( + %src: !pto.vmi.vreg<256xbf16>, + %dst: !pto.ptr, + %off: index) { + %pair = pto.vmi.vinterpret_cast %src + : !pto.vmi.vreg<256xbf16> -> !pto.vmi.vreg<128x!pto.bf16x2> + %r = pto.vmi.vcvt %pair {rounding = "A"} + : !pto.vmi.vreg<128x!pto.bf16x2> -> !pto.vmi.vreg<128x!pto.f4E2M1x2> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<128x!pto.f4E2M1x2>, !pto.ptr + return + } + + // Single register, rounding Z: 128 bf16 -> 64 bf16x2 -> 64 f4E1M2x2. + func.func @bf16x2_to_f4e1m2x2_z( + %src: !pto.vmi.vreg<128xbf16>, + %dst: !pto.ptr, + %off: index) { + %pair = pto.vmi.vinterpret_cast %src + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + %r = pto.vmi.vcvt %pair {rounding = "Z"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<64x!pto.f4E1M2x2>, !pto.ptr + return + } + + // Contract-specific rounding F must survive unified and legacy lowering. + func.func @bf16x2_to_f4e1m2x2_f( + %src: !pto.vmi.vreg<128xbf16>, + %dst: !pto.ptr, + %off: index) { + %pair = pto.vmi.vinterpret_cast %src + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + %r = pto.vmi.vcvt %pair {rounding = "F"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<64x!pto.f4E1M2x2>, !pto.ptr + return + } + + // Contract-specific rounding C on the second packed FP4 format. + func.func @bf16x2_to_f4e2m1x2_c( + %src: !pto.vmi.vreg<128xbf16>, + %dst: !pto.ptr, + %off: index) { + %pair = pto.vmi.vinterpret_cast %src + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + %r = pto.vmi.vcvt %pair {rounding = "C"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E2M1x2> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<64x!pto.f4E2M1x2>, !pto.ptr + return + } + + // Omitted rounding defaults to R for bf16x2 -> f4x2. + func.func @bf16x2_to_f4e2m1x2_default( + %src: !pto.vmi.vreg<128xbf16>, + %dst: !pto.ptr, + %off: index) { + %pair = pto.vmi.vinterpret_cast %src + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + %r = pto.vmi.vcvt %pair + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E2M1x2> + pto.vmi.vstore %r, %dst[%off] + : !pto.vmi.vreg<64x!pto.f4E2M1x2>, !pto.ptr + return + } +} + +// LOWER-LABEL: func.func @bf16x2_to_f4e1m2x2_r( +// LOWER: pto.vcvt {{.*}} {part = "P0", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER-NOT: pto.vor + +// LOWER-LABEL: func.func @bf16x2_to_f4e2m1x2_a( +// LOWER-COUNT-2: pto.vcvt {{.*}} {part = "P0", rnd = "A"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E2M1x2> +// LOWER-NOT: pto.vor + +// LOWER-LABEL: func.func @bf16x2_to_f4e1m2x2_z( +// LOWER: pto.vcvt {{.*}} {part = "P0", rnd = "Z"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER-NOT: pto.vor + +// LOWER-LABEL: func.func @bf16x2_to_f4e1m2x2_f( +// LOWER: pto.vcvt {{.*}} {part = "P0", rnd = "F"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E1M2x2> +// LOWER-NOT: pto.vor + +// LOWER-LABEL: func.func @bf16x2_to_f4e2m1x2_c( +// LOWER: pto.vcvt {{.*}} {part = "P0", rnd = "C"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E2M1x2> +// LOWER-NOT: pto.vor + +// LOWER-LABEL: func.func @bf16x2_to_f4e2m1x2_default( +// LOWER: pto.vcvt {{.*}} {part = "P0", rnd = "R"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<256x!pto.f4E2M1x2> +// LOWER-NOT: pto.vor diff --git a/test/lit/vmi_new/vmi_vcvt_bf16x2_f4x2_lane_mismatch_invalid.pto b/test/lit/vmi_new/vmi_vcvt_bf16x2_f4x2_lane_mismatch_invalid.pto new file mode 100644 index 0000000000..9ff303b2dc --- /dev/null +++ b/test/lit/vmi_new/vmi_vcvt_bf16x2_f4x2_lane_mismatch_invalid.pto @@ -0,0 +1,21 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @bf16x2_f4x2_lane_mismatch( + %source: !pto.vmi.vreg<64x!pto.bf16x2>) { + %result = pto.vmi.vcvt %source {rounding = "R"} + : !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<128x!pto.f4E1M2x2> + return + } +} + +// CHECK: requires source and result logical lane counts to match diff --git a/test/lit/vmi_new/vmi_vcvt_bf16x2_f4x2_rounding_invalid.pto b/test/lit/vmi_new/vmi_vcvt_bf16x2_f4x2_rounding_invalid.pto new file mode 100644 index 0000000000..da9c59820a --- /dev/null +++ b/test/lit/vmi_new/vmi_vcvt_bf16x2_f4x2_rounding_invalid.pto @@ -0,0 +1,67 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -split-input-file -verify-diagnostics + +// The VMI contract mirrors the physical bf16-to-f4 "RAFZC" set. H and malformed +// multi-character tokens must be rejected before reaching the VPTO verifier. + +module { + func.func @unified_h_invalid(%s: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{rounding is not valid for this fp-to-fp conversion type pair}} + %r = pto.vmi.vcvt %s {rounding = "H"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + return + } +} + +// ----- + +module { + func.func @legacy_h_invalid(%s: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{rounding attr is not valid for this fp-to-fp conversion type pair}} + %r = pto.vmi.truncf %s {rounding = "H"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E2M1x2> + return + } +} + +// ----- + +module { + func.func @unified_bad_token(%s: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{rounding is not valid for this fp-to-fp conversion type pair}} + %r = pto.vmi.vcvt %s {rounding = "Q"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + return + } +} + +// ----- + +module { + func.func @unified_multichar_invalid( + %s: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{rounding must be a single-character mode token}} + %r = pto.vmi.vcvt %s {rounding = "RA"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + return + } +} + +// ----- + +module { + func.func @legacy_multichar_invalid( + %s: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{rounding attr must be a single-character mode token}} + %r = pto.vmi.truncf %s {rounding = "RA"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E2M1x2> + return + } +} diff --git a/test/lit/vmi_new/vmi_vcvt_bf16x2_f4x2_saturate_invalid.pto b/test/lit/vmi_new/vmi_vcvt_bf16x2_f4x2_saturate_invalid.pto new file mode 100644 index 0000000000..28b52f11dd --- /dev/null +++ b/test/lit/vmi_new/vmi_vcvt_bf16x2_f4x2_saturate_invalid.pto @@ -0,0 +1,32 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -split-input-file -verify-diagnostics + +// N1: bf16x2 -> f4x2 narrows WITHOUT saturation (VPTO contract requiresSat=false). +// A saturate attribute must be rejected by the unified vcvt verifier. + +module { + func.func @bf16x2_sat_invalid(%s: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'saturate' attribute is not valid for this fp-to-fp narrow conversion (no saturation)}} + %r = pto.vmi.vcvt %s {saturate = "SAT"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + return + } +} + +// ----- + +module { + func.func @bf16x2_nosat_invalid(%s: !pto.vmi.vreg<64x!pto.bf16x2>) { + // expected-error@+1 {{'saturate' attribute is not valid for this fp-to-fp narrow conversion (no saturation)}} + %r = pto.vmi.vcvt %s {saturate = "NOSAT"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E2M1x2> + return + } +} diff --git a/test/lit/vmi_new/vmi_vcvt_f4x2_bf16x2_widen_attrs_invalid.pto b/test/lit/vmi_new/vmi_vcvt_f4x2_bf16x2_widen_attrs_invalid.pto new file mode 100644 index 0000000000..db806860e9 --- /dev/null +++ b/test/lit/vmi_new/vmi_vcvt_f4x2_bf16x2_widen_attrs_invalid.pto @@ -0,0 +1,35 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -split-input-file -verify-diagnostics + +// The f4x2 -> bf16x2 widen has no rounding or saturate semantics: rounding +// is only valid for fp-narrowing and saturate only for fp-narrow / int-narrow. +// Reject both at the unified vcvt verifier. + +module { + func.func @widen_rounding_invalid( + %s: !pto.vmi.vreg<256x!pto.f4E1M2x2>) { + // expected-error@+1 {{'rounding' attribute is only valid for fp-narrowing conversions}} + %r = pto.vmi.vcvt %s {rounding = "R"} + : !pto.vmi.vreg<256x!pto.f4E1M2x2> -> !pto.vmi.vreg<256x!pto.bf16x2> + return + } +} + +// ----- + +module { + func.func @widen_saturate_invalid( + %s: !pto.vmi.vreg<256x!pto.f4E1M2x2>) { + // expected-error@+1 {{'saturate' attribute is only valid for fp-narrow / int-narrow conversions}} + %r = pto.vmi.vcvt %s {saturate = "SAT"} + : !pto.vmi.vreg<256x!pto.f4E1M2x2> -> !pto.vmi.vreg<256x!pto.bf16x2> + return + } +} diff --git a/test/lit/vmi_new/vmi_vcvt_f4x2_to_bf16_invalid.pto b/test/lit/vmi_new/vmi_vcvt_f4x2_to_bf16_invalid.pto new file mode 100644 index 0000000000..b965de8196 --- /dev/null +++ b/test/lit/vmi_new/vmi_vcvt_f4x2_to_bf16_invalid.pto @@ -0,0 +1,24 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -verify-diagnostics + +// The f4x2 -> bf16x2 widen is supported, but the bare f4x2 -> bf16 (16-bit) +// widen is not: bf16x2 is the only legal bf16 carrier for f4 dequant. Reject +// at the unified conversion boundary instead of allowing it to fail during +// lowering. + +module { + func.func @widen_bare_bf16_unsupported( + %s: !pto.vmi.vreg<128x!pto.f4E1M2x2>) { + // expected-error@+1 {{'pto.vmi.vcvt' op unsupported packed fp-to-fp conversion element type pair}} + %r = pto.vmi.vcvt %s + : !pto.vmi.vreg<128x!pto.f4E1M2x2> -> !pto.vmi.vreg<128xbf16> + return + } +} diff --git a/test/lit/vpto/bf16x2_vreg_and_bitcast_verify.pto b/test/lit/vpto/bf16x2_vreg_and_bitcast_verify.pto new file mode 100644 index 0000000000..a93f7acea5 --- /dev/null +++ b/test/lit/vpto/bf16x2_vreg_and_bitcast_verify.pto @@ -0,0 +1,31 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not pto-test-opt %s -split-input-file 2>&1 | FileCheck %s + +module { + func.func @bf16x2_vreg_and_bitcast( + %packed: !pto.vreg<64x!pto.bf16x2>) -> !pto.vreg<64x!pto.bf16x2> { + %bf16 = pto.vbitcast %packed + : !pto.vreg<64x!pto.bf16x2> -> !pto.vreg<128xbf16> + %roundtrip = pto.vbitcast %bf16 + : !pto.vreg<128xbf16> -> !pto.vreg<64x!pto.bf16x2> + return %roundtrip : !pto.vreg<64x!pto.bf16x2> + } +} + +// ----- + +module { + func.func @bf16x2_vreg_oversized( + %packed: !pto.vreg<128x!pto.bf16x2>) { + return + } +} + +// CHECK: expected exactly 256 bytes diff --git a/test/lit/vpto/vlogic_bf16x2_verify_invalid.pto b/test/lit/vpto/vlogic_bf16x2_verify_invalid.pto new file mode 100644 index 0000000000..9a9c936a61 --- /dev/null +++ b/test/lit/vpto/vlogic_bf16x2_verify_invalid.pto @@ -0,0 +1,46 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -verify-diagnostics -split-input-file + +module { + func.func @vand_bf16x2_invalid( + %lhs: !pto.vreg<64x!pto.bf16x2>, + %rhs: !pto.vreg<64x!pto.bf16x2>, + %mask: !pto.mask) { + // expected-error @+1 {{'pto.vand' op does not support bf16x2 vector elements; low-precision bitwise operations require an 8-bit payload type}} + %out = pto.vand %lhs, %rhs, %mask : !pto.vreg<64x!pto.bf16x2>, !pto.vreg<64x!pto.bf16x2>, !pto.mask -> !pto.vreg<64x!pto.bf16x2> + return + } +} + +// ----- + +module { + func.func @vor_bf16x2_invalid( + %lhs: !pto.vreg<64x!pto.bf16x2>, + %rhs: !pto.vreg<64x!pto.bf16x2>, + %mask: !pto.mask) { + // expected-error @+1 {{'pto.vor' op does not support bf16x2 vector elements; low-precision bitwise operations require an 8-bit payload type}} + %out = pto.vor %lhs, %rhs, %mask : !pto.vreg<64x!pto.bf16x2>, !pto.vreg<64x!pto.bf16x2>, !pto.mask -> !pto.vreg<64x!pto.bf16x2> + return + } +} + +// ----- + +module { + func.func @vxor_bf16x2_invalid( + %lhs: !pto.vreg<64x!pto.bf16x2>, + %rhs: !pto.vreg<64x!pto.bf16x2>, + %mask: !pto.mask) { + // expected-error @+1 {{'pto.vxor' op does not support bf16x2 vector elements; low-precision bitwise operations require an 8-bit payload type}} + %out = pto.vxor %lhs, %rhs, %mask : !pto.vreg<64x!pto.bf16x2>, !pto.vreg<64x!pto.bf16x2>, !pto.mask -> !pto.vreg<64x!pto.bf16x2> + return + } +} diff --git a/test/lit/vpto/vmi_bf16x2_direct_load_vcvt_llvm.pto b/test/lit/vpto/vmi_bf16x2_direct_load_vcvt_llvm.pto new file mode 100644 index 0000000000..f08e1df7ee --- /dev/null +++ b/test/lit/vpto/vmi_bf16x2_direct_load_vcvt_llvm.pto @@ -0,0 +1,68 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>/dev/null | FileCheck %s --check-prefix=VPTO --implicit-check-not=pto.vor +// RUN: ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-backend=vpto --emit-vpto-llvm-ir %s -o - 2>&1 | FileCheck %s --check-prefix=LLVM + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @bf16x2_direct_load_vcvt( + %src: !pto.ptr, + %dst: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + pto.vecscope { + %wide = pto.vmi.vload %src[%c0] + : !pto.ptr + -> !pto.vmi.vreg<256x!pto.bf16x2> + %packed = pto.vmi.vcvt %wide {rounding = "R"} + : !pto.vmi.vreg<256x!pto.bf16x2> + -> !pto.vmi.vreg<256x!pto.f4E1M2x2> + pto.vmi.vstore %packed, %dst[%c0] + : !pto.vmi.vreg<256x!pto.f4E1M2x2>, + !pto.ptr + } + return + } + + func.func @bf16x2_direct_load_vcvt_e2m1( + %src: !pto.ptr, + %dst: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + pto.vecscope { + %wide = pto.vmi.vload %src[%c0] + : !pto.ptr + -> !pto.vmi.vreg<64x!pto.bf16x2> + %packed = pto.vmi.vcvt %wide {rounding = "A"} + : !pto.vmi.vreg<64x!pto.bf16x2> + -> !pto.vmi.vreg<64x!pto.f4E2M1x2> + pto.vmi.vstore %packed, %dst[%c0] + : !pto.vmi.vreg<64x!pto.f4E2M1x2>, + !pto.ptr + } + return + } +} + +// VPTO-LABEL: func.func @bf16x2_direct_load_vcvt +// VPTO-COUNT-4: pto.vlds {{.*}} : !pto.ptr -> !pto.vreg<64x!pto.bf16x2> +// VPTO-COUNT-4: pto.vcvt {{.*}} {part = "P0", rnd = "R"} +// VPTO-NOT: pto.vmi. +// VPTO-NOT: !pto.vmi. + +// LLVM-LABEL: define void @bf16x2_direct_load_vcvt_mix_aiv +// LLVM-COUNT-4: call <64 x i32> @llvm.hivm.vldsx1.v64bf16x2 +// LLVM-COUNT-4: call <256 x float4e1m2x2> @llvm.hivm.vcvtff2.bf162f4e1m2x2.x + +// VPTO-LABEL: func.func @bf16x2_direct_load_vcvt_e2m1 +// VPTO: pto.vlds {{.*}} : !pto.ptr -> !pto.vreg<64x!pto.bf16x2> +// VPTO: pto.vcvt {{.*}} {part = "P0", rnd = "A"} +// VPTO: pto.vsts +// VPTO-NOT: pto.vmi. + +// LLVM-LABEL: define void @bf16x2_direct_load_vcvt_e2m1_mix_aiv +// LLVM: call <64 x i32> @llvm.hivm.vldsx1.v64bf16x2 +// LLVM: call <256 x float4e2m1x2> @llvm.hivm.vcvtff2.bf162f4e2m1x2.x diff --git a/test/lit/vpto/vmi_f4x2_to_bf16x2_vcvt_llvm.pto b/test/lit/vpto/vmi_f4x2_to_bf16x2_vcvt_llvm.pto new file mode 100644 index 0000000000..5080f9f641 --- /dev/null +++ b/test/lit/vpto/vmi_f4x2_to_bf16x2_vcvt_llvm.pto @@ -0,0 +1,69 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>/dev/null | FileCheck %s --check-prefix=VPTO --implicit-check-not=pto.vor +// RUN: ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-backend=vpto --emit-vpto-llvm-ir %s -o - 2>&1 | FileCheck %s --check-prefix=LLVM + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @f4x2_to_bf16x2_vcvt_store( + %src: !pto.ptr, + %dst: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + pto.vecscope { + %packed = pto.vmi.vload %src[%c0] + : !pto.ptr + -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + %wide = pto.vmi.vcvt %packed + : !pto.vmi.vreg<64x!pto.f4E1M2x2> + -> !pto.vmi.vreg<64x!pto.bf16x2> + pto.vmi.vstore %wide, %dst[%c0] + : !pto.vmi.vreg<64x!pto.bf16x2>, + !pto.ptr + } + return + } + + // 256-lane d4 path: Packed4 widen produces 4 bf16x2 registers; the store + // reorders them through a vintlv tree. bf16x2 has no dedicated vintlv + // intrinsic, so the tree lowers to the <64 x i32> intrinsic (bf16x2 is a + // 32-bit packed pair at the LLVM ABI; vintlv is a bit-level lane shuffle). + func.func @f4x2_to_bf16x2_d4_vcvt_store( + %src: !pto.ptr, + %dst: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + pto.vecscope { + %packed = pto.vmi.vload %src[%c0] + : !pto.ptr + -> !pto.vmi.vreg<256x!pto.f4E1M2x2> + %wide = pto.vmi.vcvt %packed + : !pto.vmi.vreg<256x!pto.f4E1M2x2> + -> !pto.vmi.vreg<256x!pto.bf16x2> + pto.vmi.vstore %wide, %dst[%c0] + : !pto.vmi.vreg<256x!pto.bf16x2>, + !pto.ptr + } + return + } +} + +// VPTO-LABEL: func.func @f4x2_to_bf16x2_vcvt_store +// VPTO: pto.vlds {{.*}} : !pto.ptr -> !pto.vreg<256x!pto.f4E1M2x2> +// VPTO: pto.vcvt {{.*}} {part = "P0"} : !pto.vreg<256x!pto.f4E1M2x2>, !pto.mask -> !pto.vreg<128xbf16> +// VPTO: pto.vbitcast {{.*}} : !pto.vreg<128xbf16> -> !pto.vreg<64x!pto.bf16x2> +// VPTO: pto.vsts {{.*}} : !pto.vreg<64x!pto.bf16x2>, !pto.ptr, !pto.mask +// VPTO-NOT: pto.vmi. +// VPTO-NOT: !pto.vmi. + +// LLVM-LABEL: define void @f4x2_to_bf16x2_vcvt_store_mix_aiv +// LLVM: call <128 x bfloat> @llvm.hivm.vcvtff2.f4e1m2x22bf16.x(<256 x float4e1m2x2> {{.*}}, <256 x i1> {{.*}}, i32 0) +// LLVM: call void @llvm.hivm.vstsx1.v64bf16x2(<64 x i32> {{.*}}, ptr addrspace(6) {{.*}}, i32 {{.*}}, i32 2, i32 0, <256 x i1> {{.*}}) + +// LLVM-LABEL: define void @f4x2_to_bf16x2_d4_vcvt_store_mix_aiv +// LLVM: call <128 x bfloat> @llvm.hivm.vcvtff2.f4e1m2x22bf16.x(<256 x float4e1m2x2> {{.*}}, <256 x i1> {{.*}}, i32 0) +// LLVM-COUNT-4: call { <64 x i32>, <64 x i32> } @llvm.hivm.vintlv.v64i32(<64 x i32> {{.*}}, <64 x i32> {{.*}}) +// LLVM-COUNT-4: call void @llvm.hivm.vstsx1.v64bf16x2(<64 x i32> {{.*}}, ptr addrspace(6) {{.*}}, i32 {{.*}}, i32 2, i32 0, <256 x i1> {{.*}}) diff --git a/test/lit/vpto/vmi_fp4_e1_packed_surface_verify_invalid.pto b/test/lit/vpto/vmi_fp4_e1_packed_surface_verify_invalid.pto index 706ef11300..6a429b7f39 100644 --- a/test/lit/vpto/vmi_fp4_e1_packed_surface_verify_invalid.pto +++ b/test/lit/vpto/vmi_fp4_e1_packed_surface_verify_invalid.pto @@ -6,14 +6,19 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s +// Previously the packed FP4 pair types were rejected as VMI logical element +// types ("packed FP4 ... not a supported VMI surface"). Since f4x2 is now a +// first-class VMI element type (bf16x2 -> f4x2 vcvt), this is a positive test: +// the vreg verifies OK and the function round-trips. -module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { - func.func @vmi_fp4_e1_packed_surface_invalid( - %arg0: !pto.vmi.vreg<256x!pto.f4E1M2x2>) attributes {pto.kernel} { - return +// RUN: pto-test-opt %s -verify-diagnostics -split-input-file | FileCheck %s + +module { + func.func @vmi_fp4_e1_packed_surface( + %arg0: !pto.vmi.vreg<256x!pto.f4E1M2x2>) -> !pto.vmi.vreg<256x!pto.f4E1M2x2> { + return %arg0 : !pto.vmi.vreg<256x!pto.f4E1M2x2> } } -// CHECK: error: '!pto.vmi.vreg<256x!pto.f4E1M2x2>' uses a packed FP4 physical pair type as a VMI logical element type -// CHECK-SAME: packed FP4 input/output is not a supported VMI surface +// CHECK-LABEL: func.func @vmi_fp4_e1_packed_surface +// CHECK-NOT: error diff --git a/test/lit/vpto/vmi_fp4_packed_surface_verify_invalid.pto b/test/lit/vpto/vmi_fp4_packed_surface_verify_invalid.pto index 4fed3df30c..f198a4d58a 100644 --- a/test/lit/vpto/vmi_fp4_packed_surface_verify_invalid.pto +++ b/test/lit/vpto/vmi_fp4_packed_surface_verify_invalid.pto @@ -6,14 +6,19 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s +// Previously the packed FP4 pair types were rejected as VMI logical element +// types ("packed FP4 ... not a supported VMI surface"). Since f4x2 is now a +// first-class VMI element type (bf16x2 -> f4x2 vcvt), this is a positive test: +// the vreg verifies OK and the function round-trips. -module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { - func.func @vmi_fp4_packed_surface_invalid( - %arg0: !pto.vmi.vreg<256x!pto.f4E2M1x2>) attributes {pto.kernel} { - return +// RUN: pto-test-opt %s -verify-diagnostics -split-input-file | FileCheck %s + +module { + func.func @vmi_fp4_packed_surface( + %arg0: !pto.vmi.vreg<256x!pto.f4E2M1x2>) -> !pto.vmi.vreg<256x!pto.f4E2M1x2> { + return %arg0 : !pto.vmi.vreg<256x!pto.f4E2M1x2> } } -// CHECK: error: '!pto.vmi.vreg<256x!pto.f4E2M1x2>' uses a packed FP4 physical pair type as a VMI logical element type -// CHECK-SAME: packed FP4 input/output is not a supported VMI surface +// CHECK-LABEL: func.func @vmi_fp4_packed_surface +// CHECK-NOT: error diff --git a/test/python/low_precision_types.py b/test/python/low_precision_types.py index 38e9999de7..3e3f6e026c 100644 --- a/test/python/low_precision_types.py +++ b/test/python/low_precision_types.py @@ -25,12 +25,14 @@ def main() -> None: f8e8m0 = pto.F8E8M0Type.get(ctx) f4e1 = pto.F4E1M2x2Type.get(ctx) f4e2 = pto.F4E2M1x2Type.get(ctx) + bf16x2 = pto.BF16x2Type.get(ctx) assert_contains(str(hif8), "hif8") assert_contains(str(hif8x2), "hif8x2") assert_contains(str(f8e8m0), "f8E8M0") assert_contains(str(f4e1), "f4E1M2x2") assert_contains(str(f4e2), "f4E2M1x2") + assert_contains(str(bf16x2), "bf16x2") print("low_precision_types: PASS") diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/compare.py b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/compare.py new file mode 100644 index 0000000000..d3a888e0f3 --- /dev/null +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/compare.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import sys + +import numpy as np + + +def main() -> None: + golden = np.fromfile("golden_v2.bin", dtype=np.uint16) + output = np.fromfile("v2.bin", dtype=np.uint16) + if golden.shape != output.shape: + print(f"[ERROR] size mismatch golden={golden.size} output={output.size}") + sys.exit(2) + if not np.array_equal(golden, output): + diff = np.nonzero(golden != output)[0] + idx = int(diff[0]) + print( + f"[ERROR] compare failed idx={idx} " + f"golden=0x{int(golden[idx]):04x} " + f"output=0x{int(output[idx]):04x}" + ) + sys.exit(2) + print("[INFO] compare passed") + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/golden.py b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/golden.py new file mode 100644 index 0000000000..0babbced04 --- /dev/null +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/golden.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# Reference for f4E1M2x2 -> bf16x2 (contiguous path). +# Each f4E1M2x2 byte packs two 4-bit FP4 E1M2 codes: low nibble = even bf16 +# lane, high nibble = odd bf16 lane (mirrors the quant case PAIR_HI_FIRST). +# Each code dequantizes to one bf16 value; the pair becomes one bf16x2. +# +# Calibration note: the E1M2 table and nibble order below are provisional +# oracles mirroring quant-bf16x2-to-f4x2-contiguous. Confirm them on the first +# A5 NPU run before treating a strict golden mismatch as a compiler regression. + +import argparse +from pathlib import Path + +import numpy as np + +SRC_ELEMS = 256 # f4x2 input bytes +DST_ELEMS = 512 # bf16 output lanes + +# Magnitude values for E1M2 code 0..7 (code bit3 = sign). IEEE-style table. +E1M2_VALUES = [0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75] +PAIR_HI_FIRST = True # True: byte = (f4(odd) << 4) | f4(even); False: reversed + +# Input f4 codes (all 16 codes, values exactly representable in the table). +CODES = np.array( + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], dtype=np.uint8 +) + + +def f32_to_bf16_bits(values: np.ndarray) -> np.ndarray: + """Round f32 to nearest bf16 and return its 16-bit pattern.""" + f32 = values.astype(np.float32) + bits = f32.view(np.uint32) + # round-to-nearest-even on the low 16 bits + lsb = (bits >> 16) & 1 + rounding = 0x7FFF + lsb + bf16 = (bits + rounding) >> 16 + return bf16.astype(np.uint16) + + +def dequantize_f4e1m2(codes: np.ndarray) -> np.ndarray: + """f4E1M2 code (0..15) -> f32 value.""" + sign = ((codes >> 3) & 1).astype(np.int32) + mag = np.array(E1M2_VALUES, dtype=np.float32)[(codes & 0x7).astype(np.int32)] + values = np.where(sign == 1, -mag, mag).astype(np.float32) + return values + + +def generate(output_dir: Path) -> None: + # SRC_ELEMS f4x2 bytes pack 2 * SRC_ELEMS f4 codes. + num_codes = SRC_ELEMS * 2 + repeats = (num_codes + len(CODES) - 1) // len(CODES) + codes = np.tile(CODES, repeats)[:num_codes].astype(np.uint8) + + # Pack into f4x2 bytes: low nibble = even lane code, high nibble = odd. + even_codes = codes[0::2] + odd_codes = codes[1::2] + assert even_codes.size == odd_codes.size == SRC_ELEMS + if PAIR_HI_FIRST: + src_bytes = ((odd_codes.astype(np.uint8) << 4) | even_codes).astype(np.uint8) + else: + src_bytes = ((even_codes.astype(np.uint8) << 4) | odd_codes).astype(np.uint8) + + # Dequantize: bf16[2j] = low nibble, bf16[2j+1] = high nibble. + vals_even = dequantize_f4e1m2(even_codes) + vals_odd = dequantize_f4e1m2(odd_codes) + out_f32 = np.empty(DST_ELEMS, dtype=np.float32) + out_f32[0::2] = vals_even + out_f32[1::2] = vals_odd + dst_bf16 = f32_to_bf16_bits(out_f32) + + src_buf = src_bytes.astype(np.uint8) + dst_buf = np.full(DST_ELEMS, 0xA5, dtype=np.uint16) + + output_dir.mkdir(parents=True, exist_ok=True) + src_buf.tofile(output_dir / "v1.bin") + dst_buf.tofile(output_dir / "v2.bin") + dst_bf16.tofile(output_dir / "golden_v2.bin") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, default=Path(".")) + args = parser.parse_args() + generate(args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/kernel.pto b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/kernel.pto new file mode 100644 index 0000000000..9cddb046bb --- /dev/null +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/kernel.pto @@ -0,0 +1,50 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// ----------------------------------------------------------------------------- +// case: vmi_new/dequant-f4x2-to-bf16x2-contiguous +// family: conversion +// target_ops: pto.vmi.vcvt (f4E1M2x2 -> bf16x2) +// scenarios: f4x2 contiguous load -> vcvt (widen, Packed4 P0-P3) -> bf16x2 d4 store +// ----------------------------------------------------------------------------- +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_dequant_f4x2_to_bf16x2_contiguous_kernel(%src_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c256_i64 = arith.constant 256 : i64 + %c1024_i64 = arith.constant 1024 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_dst_bf16 = pto.castptr %c1024_i64 : i64 -> !pto.ptr + %ub_dst = pto.castptr %c1024_i64 : i64 -> !pto.ptr + + // src: 256 f4x2 = 256 bytes; dst: 256 bf16x2 = 1024 bytes. + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c256_i64 + nburst(%c1_i64, %c256_i64, %c256_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %dst_gm, %ub_dst_bf16, %c0_i64, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.vecscope { + %packed = pto.vmi.vload %ub_src[%c0] : !pto.ptr -> !pto.vmi.vreg<256x!pto.f4E1M2x2> + %wide = pto.vmi.vcvt %packed + : !pto.vmi.vreg<256x!pto.f4E1M2x2> -> !pto.vmi.vreg<256x!pto.bf16x2> + pto.vmi.vstore %wide, %ub_dst[%c0] + : !pto.vmi.vreg<256x!pto.bf16x2>, !pto.ptr + } + + pto.mte_ub_gm %ub_dst_bf16, %dst_gm, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/launch.cpp b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/launch.cpp new file mode 100644 index 0000000000..72d317fcad --- /dev/null +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/launch.cpp @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef __VEC_SCOPE__ +#define __VEC_SCOPE__ +#endif +#if defined(__CCE_AICORE__) && defined(__NPU_ARCH__) && (__NPU_ARCH__ == 2201) +typedef struct { unsigned char v; } hifloat8_t; +typedef struct { unsigned char v; } float8_e4m3_t; +typedef struct { unsigned char v; } float8_e5m2_t; +typedef struct { unsigned char v; } float8_e8m0_t; +typedef struct { unsigned char v; } float4_e1m2x2_t; +typedef struct { unsigned char v; } float4_e2m1x2_t; +#endif +#include +#if !defined(__CCE_AICORE__) && !defined(TMRGSORT_HPP) +struct MrgSortExecutedNumList { + uint16_t mrgSortList0; + uint16_t mrgSortList1; + uint16_t mrgSortList2; + uint16_t mrgSortList3; +}; +#endif +#ifndef __CPU_SIM +#include "acl/acl.h" +#endif + +extern "C" __global__ [aicore] void +vmi_dequant_f4x2_to_bf16x2_contiguous_kernel(__gm__ uint8_t *src, __gm__ uint16_t *dst); + +void LaunchVmi_dequant_f4x2_to_bf16x2_contiguous_kernel(uint8_t *src, uint16_t *dst, + void *stream) { + vmi_dequant_f4x2_to_bf16x2_contiguous_kernel<<<1, nullptr, stream>>>( + (__gm__ uint8_t *)src, (__gm__ uint16_t *)dst); +} diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/main.cpp b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/main.cpp new file mode 100644 index 0000000000..fbcc940774 --- /dev/null +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/main.cpp @@ -0,0 +1,79 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include + +using namespace PtoTestCommon; + +#define ACL_CHECK(expr) \ + do { \ + const aclError _ret = (expr); \ + if (_ret != ACL_SUCCESS) { \ + std::fprintf(stderr, "[ERROR] %s failed: %d (%s:%d)\n", #expr, \ + (int)_ret, __FILE__, __LINE__); \ + rc = 1; \ + goto cleanup; \ + } \ + } while (0) + +void LaunchVmi_dequant_f4x2_to_bf16x2_contiguous_kernel(uint8_t *src, uint16_t *dst, + void *stream); + +int main() { + constexpr size_t kSrcElems = 256; // f4x2 (1 byte per pair) + constexpr size_t kDstElems = 512; // bf16 output lanes + size_t srcBytes = kSrcElems * sizeof(uint8_t); + size_t dstBytes = kDstElems * sizeof(uint16_t); + uint8_t *srcHost = nullptr; + uint8_t *srcDevice = nullptr; + uint16_t *dstHost = nullptr; + uint16_t *dstDevice = nullptr; + int rc = 0; + bool aclInited = false; + bool deviceSet = false; + int deviceId = 0; + aclrtStream stream = nullptr; + + ACL_CHECK(aclInit(nullptr)); + aclInited = true; + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) + deviceId = std::atoi(envDevice); + ACL_CHECK(aclrtSetDevice(deviceId)); + deviceSet = true; + ACL_CHECK(aclrtCreateStream(&stream)); + ACL_CHECK(aclrtMallocHost((void **)(&srcHost), srcBytes)); + ACL_CHECK(aclrtMallocHost((void **)(&dstHost), dstBytes)); + ACL_CHECK(aclrtMalloc((void **)&srcDevice, srcBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + ACL_CHECK(aclrtMalloc((void **)&dstDevice, dstBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + + ReadFile("./v1.bin", srcBytes, srcHost, srcBytes); + ReadFile("./v2.bin", dstBytes, dstHost, dstBytes); + ACL_CHECK(aclrtMemcpy(srcDevice, srcBytes, srcHost, srcBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + ACL_CHECK(aclrtMemcpy(dstDevice, dstBytes, dstHost, dstBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + LaunchVmi_dequant_f4x2_to_bf16x2_contiguous_kernel(srcDevice, dstDevice, stream); + ACL_CHECK(aclrtSynchronizeStream(stream)); + ACL_CHECK(aclrtMemcpy(dstHost, dstBytes, dstDevice, dstBytes, ACL_MEMCPY_DEVICE_TO_HOST)); + WriteFile("./v2.bin", dstHost, dstBytes); + +cleanup: + aclrtFree(srcDevice); + aclrtFree(dstDevice); + aclrtFreeHost(srcHost); + aclrtFreeHost(dstHost); + if (stream) + aclrtDestroyStream(stream); + if (deviceSet) + aclrtResetDevice(deviceId); + if (aclInited) + aclFinalize(); + return rc; +} diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/ptoas.flags b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/ptoas.flags new file mode 100644 index 0000000000..5d9dc67120 --- /dev/null +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/ptoas.flags @@ -0,0 +1 @@ +--pto-arch a5 --pto-backend=vpto diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/compare.py b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/compare.py new file mode 100644 index 0000000000..d3a888e0f3 --- /dev/null +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/compare.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import sys + +import numpy as np + + +def main() -> None: + golden = np.fromfile("golden_v2.bin", dtype=np.uint16) + output = np.fromfile("v2.bin", dtype=np.uint16) + if golden.shape != output.shape: + print(f"[ERROR] size mismatch golden={golden.size} output={output.size}") + sys.exit(2) + if not np.array_equal(golden, output): + diff = np.nonzero(golden != output)[0] + idx = int(diff[0]) + print( + f"[ERROR] compare failed idx={idx} " + f"golden=0x{int(golden[idx]):04x} " + f"output=0x{int(output[idx]):04x}" + ) + sys.exit(2) + print("[INFO] compare passed") + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/golden.py b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/golden.py new file mode 100644 index 0000000000..1ffcd05d3d --- /dev/null +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/golden.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# Reference for f4E1M2x2 -> bf16x2 (tail loop path). Mirrors +# dequant-f4x2-to-bf16x2-contiguous golden, but only 500 of the 512 lanes are +# active: iteration 4 writes 116 lanes under a dynamic mask. Lanes 500..511 +# keep the 0xA5 prefill. + +import argparse +from pathlib import Path + +import numpy as np + +SRC_ELEMS = 256 # f4x2 input bytes in the UB buffer +ACTIVE_ELEMS = 250 # logical lanes actually processed +DST_ELEMS = 512 # bf16 output lanes in the UB buffer + +# Magnitude values for E1M2 code 0..7 (code bit3 = sign). IEEE-style table. +E1M2_VALUES = [0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75] +PAIR_HI_FIRST = True # True: byte = (f4(odd) << 4) | f4(even); False: reversed + +# Input f4 codes (all 16 codes, values exactly representable in the table). +CODES = np.array( + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], dtype=np.uint8 +) + + +def f32_to_bf16_bits(values: np.ndarray) -> np.ndarray: + """Round f32 to nearest bf16 and return its 16-bit pattern.""" + f32 = values.astype(np.float32) + bits = f32.view(np.uint32) + # round-to-nearest-even on the low 16 bits + lsb = (bits >> 16) & 1 + rounding = 0x7FFF + lsb + bf16 = (bits + rounding) >> 16 + return bf16.astype(np.uint16) + + +def dequantize_f4e1m2(codes: np.ndarray) -> np.ndarray: + """f4E1M2 code (0..15) -> f32 value.""" + sign = ((codes >> 3) & 1).astype(np.int32) + mag = np.array(E1M2_VALUES, dtype=np.float32)[(codes & 0x7).astype(np.int32)] + values = np.where(sign == 1, -mag, mag).astype(np.float32) + return values + + +def generate(output_dir: Path) -> None: + # Active lanes only (ACTIVE_ELEMS f4x2 bytes = 2 * ACTIVE_ELEMS f4 codes); + # the remaining bytes stay as prefill. + num_codes = ACTIVE_ELEMS * 2 + repeats = (num_codes + len(CODES) - 1) // len(CODES) + codes = np.tile(CODES, repeats)[:num_codes].astype(np.uint8) + + # Pack into f4x2 bytes: low nibble = even lane code, high nibble = odd. + even_codes = codes[0::2] + odd_codes = codes[1::2] + assert even_codes.size == odd_codes.size == ACTIVE_ELEMS + if PAIR_HI_FIRST: + packed = ((odd_codes.astype(np.uint8) << 4) | even_codes).astype(np.uint8) + else: + packed = ((even_codes.astype(np.uint8) << 4) | odd_codes).astype(np.uint8) + + src_bytes = np.full(SRC_ELEMS, 0xA5, dtype=np.uint8) + src_bytes[:packed.size] = packed + + # Dequantize: bf16[2j] = low nibble, bf16[2j+1] = high nibble. + vals_even = dequantize_f4e1m2(even_codes) + vals_odd = dequantize_f4e1m2(odd_codes) + out_f32 = np.empty(ACTIVE_ELEMS * 2, dtype=np.float32) + out_f32[0::2] = vals_even + out_f32[1::2] = vals_odd + active_bf16 = f32_to_bf16_bits(out_f32) + + dst_bf16 = np.full(DST_ELEMS, 0xA5, dtype=np.uint16) + dst_bf16[:active_bf16.size] = active_bf16 + + dst_buf = np.full(DST_ELEMS, 0xA5, dtype=np.uint16) + + output_dir.mkdir(parents=True, exist_ok=True) + src_bytes.tofile(output_dir / "v1.bin") + dst_buf.tofile(output_dir / "v2.bin") + dst_bf16.tofile(output_dir / "golden_v2.bin") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, default=Path(".")) + args = parser.parse_args() + generate(args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/kernel.pto b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/kernel.pto new file mode 100644 index 0000000000..8c2ce1aaf4 --- /dev/null +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/kernel.pto @@ -0,0 +1,59 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// ----------------------------------------------------------------------------- +// case: vmi_new/dequant-f4x2-to-bf16x2-tail +// family: conversion +// target_ops: pto.vmi.vcvt (f4E1M2x2 -> bf16x2), masked store in a loop +// scenarios: f4x2 loop + dynamic mask tail dequant (64 lanes per iteration) +// ----------------------------------------------------------------------------- +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_dequant_f4x2_to_bf16x2_tail_kernel(%src_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c64 = arith.constant 64 : index + %c250 = arith.constant 250 : index + %c256 = arith.constant 256 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c256_i64 = arith.constant 256 : i64 + %c1024_i64 = arith.constant 1024 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_dst_bf16 = pto.castptr %c1024_i64 : i64 -> !pto.ptr + %ub_dst = pto.castptr %c1024_i64 : i64 -> !pto.ptr + + // src buffer: 256 f4x2 = 256 bytes (250 logical); dst: 256 bf16x2 = 1024 bytes. + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c256_i64 + nburst(%c1_i64, %c256_i64, %c256_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %dst_gm, %ub_dst_bf16, %c0_i64, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.vecscope { + %_:1 = scf.for %offset = %c0 to %c256 step %c64 iter_args(%rem = %c250) -> (index) { + %mask = pto.vmi.create_mask %rem : index -> !pto.vmi.mask<64xpred> + %packed = pto.vmi.vload %ub_src[%offset] : !pto.ptr -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + %wide = pto.vmi.vcvt %packed + : !pto.vmi.vreg<64x!pto.f4E1M2x2> -> !pto.vmi.vreg<64x!pto.bf16x2> + pto.vmi.vstore %wide, %ub_dst[%offset], %mask + : !pto.vmi.vreg<64x!pto.bf16x2>, !pto.ptr, + !pto.vmi.mask<64xpred> + %next = arith.subi %rem, %c64 : index + scf.yield %next : index + } + } + + pto.mte_ub_gm %ub_dst_bf16, %dst_gm, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/launch.cpp b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/launch.cpp new file mode 100644 index 0000000000..0b0ef1b7c0 --- /dev/null +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/launch.cpp @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef __VEC_SCOPE__ +#define __VEC_SCOPE__ +#endif +#if defined(__CCE_AICORE__) && defined(__NPU_ARCH__) && (__NPU_ARCH__ == 2201) +typedef struct { unsigned char v; } hifloat8_t; +typedef struct { unsigned char v; } float8_e4m3_t; +typedef struct { unsigned char v; } float8_e5m2_t; +typedef struct { unsigned char v; } float8_e8m0_t; +typedef struct { unsigned char v; } float4_e1m2x2_t; +typedef struct { unsigned char v; } float4_e2m1x2_t; +#endif +#include +#if !defined(__CCE_AICORE__) && !defined(TMRGSORT_HPP) +struct MrgSortExecutedNumList { + uint16_t mrgSortList0; + uint16_t mrgSortList1; + uint16_t mrgSortList2; + uint16_t mrgSortList3; +}; +#endif +#ifndef __CPU_SIM +#include "acl/acl.h" +#endif + +extern "C" __global__ [aicore] void +vmi_dequant_f4x2_to_bf16x2_tail_kernel(__gm__ uint8_t *src, __gm__ uint16_t *dst); + +void LaunchVmi_dequant_f4x2_to_bf16x2_tail_kernel(uint8_t *src, uint16_t *dst, + void *stream) { + vmi_dequant_f4x2_to_bf16x2_tail_kernel<<<1, nullptr, stream>>>( + (__gm__ uint8_t *)src, (__gm__ uint16_t *)dst); +} diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/main.cpp b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/main.cpp new file mode 100644 index 0000000000..2837a64e90 --- /dev/null +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/main.cpp @@ -0,0 +1,79 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include + +using namespace PtoTestCommon; + +#define ACL_CHECK(expr) \ + do { \ + const aclError _ret = (expr); \ + if (_ret != ACL_SUCCESS) { \ + std::fprintf(stderr, "[ERROR] %s failed: %d (%s:%d)\n", #expr, \ + (int)_ret, __FILE__, __LINE__); \ + rc = 1; \ + goto cleanup; \ + } \ + } while (0) + +void LaunchVmi_dequant_f4x2_to_bf16x2_tail_kernel(uint8_t *src, uint16_t *dst, + void *stream); + +int main() { + constexpr size_t kSrcElems = 256; // f4x2 bytes in the UB buffer + constexpr size_t kDstElems = 512; // bf16 output lanes in the UB buffer + size_t srcBytes = kSrcElems * sizeof(uint8_t); + size_t dstBytes = kDstElems * sizeof(uint16_t); + uint8_t *srcHost = nullptr; + uint8_t *srcDevice = nullptr; + uint16_t *dstHost = nullptr; + uint16_t *dstDevice = nullptr; + int rc = 0; + bool aclInited = false; + bool deviceSet = false; + int deviceId = 0; + aclrtStream stream = nullptr; + + ACL_CHECK(aclInit(nullptr)); + aclInited = true; + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) + deviceId = std::atoi(envDevice); + ACL_CHECK(aclrtSetDevice(deviceId)); + deviceSet = true; + ACL_CHECK(aclrtCreateStream(&stream)); + ACL_CHECK(aclrtMallocHost((void **)(&srcHost), srcBytes)); + ACL_CHECK(aclrtMallocHost((void **)(&dstHost), dstBytes)); + ACL_CHECK(aclrtMalloc((void **)&srcDevice, srcBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + ACL_CHECK(aclrtMalloc((void **)&dstDevice, dstBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + + ReadFile("./v1.bin", srcBytes, srcHost, srcBytes); + ReadFile("./v2.bin", dstBytes, dstHost, dstBytes); + ACL_CHECK(aclrtMemcpy(srcDevice, srcBytes, srcHost, srcBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + ACL_CHECK(aclrtMemcpy(dstDevice, dstBytes, dstHost, dstBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + LaunchVmi_dequant_f4x2_to_bf16x2_tail_kernel(srcDevice, dstDevice, stream); + ACL_CHECK(aclrtSynchronizeStream(stream)); + ACL_CHECK(aclrtMemcpy(dstHost, dstBytes, dstDevice, dstBytes, ACL_MEMCPY_DEVICE_TO_HOST)); + WriteFile("./v2.bin", dstHost, dstBytes); + +cleanup: + aclrtFree(srcDevice); + aclrtFree(dstDevice); + aclrtFreeHost(srcHost); + aclrtFreeHost(dstHost); + if (stream) + aclrtDestroyStream(stream); + if (deviceSet) + aclrtResetDevice(deviceId); + if (aclInited) + aclFinalize(); + return rc; +} diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/ptoas.flags b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/ptoas.flags new file mode 100644 index 0000000000..5d9dc67120 --- /dev/null +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/ptoas.flags @@ -0,0 +1 @@ +--pto-arch a5 --pto-backend=vpto diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/compare.py b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/compare.py new file mode 100644 index 0000000000..d95d3fec3c --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/compare.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import sys + +import numpy as np + + +def main() -> None: + golden = np.fromfile("golden_v2.bin", dtype=np.uint8) + output = np.fromfile("v2.bin", dtype=np.uint8) + if golden.shape != output.shape: + print(f"[ERROR] size mismatch golden={golden.size} output={output.size}") + sys.exit(2) + if not np.array_equal(golden, output): + diff = np.nonzero(golden != output)[0] + idx = int(diff[0]) + print( + f"[ERROR] compare failed idx={idx} " + f"golden=0x{int(golden[idx]):02x} " + f"output=0x{int(output[idx]):02x}" + ) + sys.exit(2) + print("[INFO] compare passed") + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/golden.py b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/golden.py new file mode 100644 index 0000000000..470f4af9d7 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/golden.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# Reference for bf16x2 -> f4E2M1x2 (second FP4 format). +# E2M1 magnitude table {0, 0.5, 1, 1.5, 2, 3, 4, 6} for codes 0..7; code bit3 +# is sign. Input values are exactly representable so no rounding ambiguity. +# Calibration note: the E2M1 table, rounding interpretation, and low/high +# nibble order below are provisional oracles. Confirm them independently on the +# first A5 NPU run before treating a strict golden mismatch as a compiler bug. + +import argparse +from pathlib import Path + +import numpy as np + +SRC_ELEMS = 128 # bf16 input lanes +DST_ELEMS = 64 # f4x2 output bytes + +# Magnitude values for E2M1 code 0..7 (code bit3 = sign). +E2M1_VALUES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] +PAIR_HI_FIRST = True + +# Exactly representable E2M1 values (both signs, spanning the range). +VALUES = np.array([0.0, 0.5, -0.5, 1.0, -1.0, 1.5, -1.5, 2.0, -2.0, 3.0, 4.0, 6.0], + dtype=np.float32) + + +def f32_to_bf16_bits(values: np.ndarray) -> np.ndarray: + """Round f32 to nearest bf16 and return its 16-bit pattern.""" + f32 = values.astype(np.float32) + bits = f32.view(np.uint32) + lsb = (bits >> 16) & 1 + rounding = 0x7FFF + lsb + bf16 = (bits + rounding) >> 16 + return bf16.astype(np.uint16) + + +def bf16_to_f32(bits: np.ndarray) -> np.ndarray: + return (bits.astype(np.uint32) << 16).view(np.float32) + + +def quantize_f4e2m1(values: np.ndarray) -> np.ndarray: + """Nearest f4E2M1 code (0..15) for each f32 value.""" + sign = (values < 0).astype(np.int32) + mag = np.abs(values) + table = np.array(E2M1_VALUES, dtype=np.float32) + idx = np.argmin(np.abs(table[None, :] - mag[:, None]), axis=1) + return ((sign << 3) | idx).astype(np.uint8) + + +def generate(output_dir: Path) -> None: + repeats = (SRC_ELEMS + len(VALUES) - 1) // len(VALUES) + src_f32 = np.tile(VALUES, repeats)[:SRC_ELEMS].astype(np.float32) + src_bf16 = f32_to_bf16_bits(src_f32) + + f4 = quantize_f4e2m1(bf16_to_f32(src_bf16)).astype(np.int32) + even = f4[0::2] + odd = f4[1::2] + if PAIR_HI_FIRST: + dst = ((odd << 4) | even).astype(np.uint8) + else: + dst = ((even << 4) | odd).astype(np.uint8) + assert dst.size == DST_ELEMS + + src_buf = src_bf16.astype(" None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, default=Path(".")) + args = parser.parse_args() + generate(args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/kernel.pto b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/kernel.pto new file mode 100644 index 0000000000..ce780d3c93 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/kernel.pto @@ -0,0 +1,59 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// ----------------------------------------------------------------------------- +// case: vmi_new/quant-bf16x2-to-f4e2m1x2 +// family: conversion +// target_ops: pto.vmi.vcvt (bf16x2 -> f4E2M1x2) +// scenarios: the SECOND FP4 format. E2M1 value table {0,0.5,1,1.5,2,3,4,6} +// differs from E1M2; probes the assumed table pending NPU calibration. +// bf16 contiguous load -> vinterpret_cast pair -> vcvt -> store. +// ----------------------------------------------------------------------------- +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_quant_bf16x2_to_f4e2m1x2_kernel(%src_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c64_i64 = arith.constant 64 : i64 + %c256_i64 = arith.constant 256 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_dst_u8 = pto.castptr %c256_i64 : i64 -> !pto.ptr + %ub_dst_f4 = pto.castptr %c256_i64 : i64 -> !pto.ptr + + // src: 128 bf16 = 256 bytes; dst: 64 f4x2 = 64 bytes. + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c256_i64 + nburst(%c1_i64, %c256_i64, %c256_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %dst_gm, %ub_dst_u8, %c0_i64, %c64_i64 + nburst(%c1_i64, %c64_i64, %c64_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + pto.vecscope { + %wide = pto.vmi.vload %ub_src[%c0] : !pto.ptr -> !pto.vmi.vreg<128xbf16> + %pair = pto.vmi.vinterpret_cast %wide + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + %packed = pto.vmi.vcvt %pair {rounding = "R"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E2M1x2> + pto.vmi.vstore %packed, %ub_dst_f4[%c0] + : !pto.vmi.vreg<64x!pto.f4E2M1x2>, !pto.ptr + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.mte_ub_gm %ub_dst_u8, %dst_gm, %c64_i64 + nburst(%c1_i64, %c64_i64, %c64_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/launch.cpp b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/launch.cpp new file mode 100644 index 0000000000..fbf989e342 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/launch.cpp @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef __VEC_SCOPE__ +#define __VEC_SCOPE__ +#endif +#if defined(__CCE_AICORE__) && defined(__NPU_ARCH__) && (__NPU_ARCH__ == 2201) +typedef struct { unsigned char v; } hifloat8_t; +typedef struct { unsigned char v; } float8_e4m3_t; +typedef struct { unsigned char v; } float8_e5m2_t; +typedef struct { unsigned char v; } float8_e8m0_t; +typedef struct { unsigned char v; } float4_e1m2x2_t; +typedef struct { unsigned char v; } float4_e2m1x2_t; +#endif +#include +#if !defined(__CCE_AICORE__) && !defined(TMRGSORT_HPP) +struct MrgSortExecutedNumList { + uint16_t mrgSortList0; + uint16_t mrgSortList1; + uint16_t mrgSortList2; + uint16_t mrgSortList3; +}; +#endif +#ifndef __CPU_SIM +#include "acl/acl.h" +#endif + +extern "C" __global__ [aicore] void +vmi_quant_bf16x2_to_f4e2m1x2_kernel(__gm__ uint16_t *src, __gm__ uint8_t *dst); + +void LaunchVmi_quant_bf16x2_to_f4e2m1x2_kernel(uint16_t *src, uint8_t *dst, + void *stream) { + vmi_quant_bf16x2_to_f4e2m1x2_kernel<<<1, nullptr, stream>>>( + (__gm__ uint16_t *)src, (__gm__ uint8_t *)dst); +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/main.cpp b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/main.cpp new file mode 100644 index 0000000000..cdd00da2a9 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/main.cpp @@ -0,0 +1,79 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include + +using namespace PtoTestCommon; + +#define ACL_CHECK(expr) \ + do { \ + const aclError _ret = (expr); \ + if (_ret != ACL_SUCCESS) { \ + std::fprintf(stderr, "[ERROR] %s failed: %d (%s:%d)\n", #expr, \ + (int)_ret, __FILE__, __LINE__); \ + rc = 1; \ + goto cleanup; \ + } \ + } while (0) + +void LaunchVmi_quant_bf16x2_to_f4e2m1x2_kernel(uint16_t *src, uint8_t *dst, + void *stream); + +int main() { + constexpr size_t kSrcElems = 128; // bf16 + constexpr size_t kDstElems = 64; // f4E2M1x2 (1 byte per pair) + size_t srcBytes = kSrcElems * sizeof(uint16_t); + size_t dstBytes = kDstElems * sizeof(uint8_t); + uint16_t *srcHost = nullptr; + uint16_t *srcDevice = nullptr; + uint8_t *dstHost = nullptr; + uint8_t *dstDevice = nullptr; + int rc = 0; + bool aclInited = false; + bool deviceSet = false; + int deviceId = 0; + aclrtStream stream = nullptr; + + ACL_CHECK(aclInit(nullptr)); + aclInited = true; + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) + deviceId = std::atoi(envDevice); + ACL_CHECK(aclrtSetDevice(deviceId)); + deviceSet = true; + ACL_CHECK(aclrtCreateStream(&stream)); + ACL_CHECK(aclrtMallocHost((void **)(&srcHost), srcBytes)); + ACL_CHECK(aclrtMallocHost((void **)(&dstHost), dstBytes)); + ACL_CHECK(aclrtMalloc((void **)&srcDevice, srcBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + ACL_CHECK(aclrtMalloc((void **)&dstDevice, dstBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + + ReadFile("./v1.bin", srcBytes, srcHost, srcBytes); + ReadFile("./v2.bin", dstBytes, dstHost, dstBytes); + ACL_CHECK(aclrtMemcpy(srcDevice, srcBytes, srcHost, srcBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + ACL_CHECK(aclrtMemcpy(dstDevice, dstBytes, dstHost, dstBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + LaunchVmi_quant_bf16x2_to_f4e2m1x2_kernel(srcDevice, dstDevice, stream); + ACL_CHECK(aclrtSynchronizeStream(stream)); + ACL_CHECK(aclrtMemcpy(dstHost, dstBytes, dstDevice, dstBytes, ACL_MEMCPY_DEVICE_TO_HOST)); + WriteFile("./v2.bin", dstHost, dstBytes); + +cleanup: + aclrtFree(srcDevice); + aclrtFree(dstDevice); + aclrtFreeHost(srcHost); + aclrtFreeHost(dstHost); + if (stream) + aclrtDestroyStream(stream); + if (deviceSet) + aclrtResetDevice(deviceId); + if (aclInited) + aclFinalize(); + return rc; +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/ptoas.flags b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/ptoas.flags new file mode 100644 index 0000000000..5d9dc67120 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4e2m1x2/ptoas.flags @@ -0,0 +1 @@ +--pto-arch a5 --pto-backend=vpto diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/compare.py b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/compare.py new file mode 100644 index 0000000000..d95d3fec3c --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/compare.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import sys + +import numpy as np + + +def main() -> None: + golden = np.fromfile("golden_v2.bin", dtype=np.uint8) + output = np.fromfile("v2.bin", dtype=np.uint8) + if golden.shape != output.shape: + print(f"[ERROR] size mismatch golden={golden.size} output={output.size}") + sys.exit(2) + if not np.array_equal(golden, output): + diff = np.nonzero(golden != output)[0] + idx = int(diff[0]) + print( + f"[ERROR] compare failed idx={idx} " + f"golden=0x{int(golden[idx]):02x} " + f"output=0x{int(output[idx]):02x}" + ) + sys.exit(2) + print("[INFO] compare passed") + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/golden.py b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/golden.py new file mode 100644 index 0000000000..d3e400e708 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/golden.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# Reference for bf16x2 -> f4E1M2x2 (contiguous path). +# Each bf16 (2-byte) quantizes to one 4-bit FP4 E1M2 code; a bf16x2 pair of +# {bf16[2j], bf16[2j+1]} produces one f4E1M2x2 byte. +# +# Calibration note: the E1M2 table, rounding interpretation, and low/high +# nibble order below are provisional oracles. Confirm them on the first A5 NPU +# run before treating a strict golden mismatch as a compiler regression. + +import argparse +from pathlib import Path + +import numpy as np + +SRC_ELEMS = 256 # bf16 input lanes +DST_ELEMS = 128 # f4x2 output bytes + +# Magnitude values for E1M2 code 0..7 (code bit3 = sign). IEEE-style table. +E1M2_VALUES = [0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75] +PAIR_HI_FIRST = True # True: byte = (f4(odd) << 4) | f4(even); False: reversed + +# Input values (exactly representable in the IEEE-style E1M2 table above). +VALUES = np.array([0.0, 1.0, -1.0, 0.5, -0.5, 1.5, -1.5, 0.25], dtype=np.float32) + + +def f32_to_bf16_bits(values: np.ndarray) -> np.ndarray: + """Round f32 to nearest bf16 and return its 16-bit pattern.""" + f32 = values.astype(np.float32) + bits = f32.view(np.uint32) + # round-to-nearest-even on the low 16 bits + lsb = (bits >> 16) & 1 + rounding = 0x7FFF + lsb + bf16 = (bits + rounding) >> 16 + return bf16.astype(np.uint16) + + +def bf16_to_f32(bits: np.ndarray) -> np.ndarray: + f32 = bits.astype(np.uint32) << 16 + return f32.view(np.float32) + + +def quantize_f4e1m2(values: np.ndarray) -> np.ndarray: + """Nearest f4E1M2 code (0..15) for each f32 value.""" + sign = (values < 0).astype(np.int32) + mag = np.abs(values) + table = np.array(E1M2_VALUES, dtype=np.float32) + idx = np.argmin(np.abs(table[None, :] - mag[:, None]), axis=1) + return ((sign << 3) | idx).astype(np.uint8) + + +def generate(output_dir: Path) -> None: + repeats = (SRC_ELEMS + len(VALUES) - 1) // len(VALUES) + src_f32 = np.tile(VALUES, repeats)[:SRC_ELEMS].astype(np.float32) + src_bf16 = f32_to_bf16_bits(src_f32) + + f4 = quantize_f4e1m2(bf16_to_f32(src_bf16)).astype(np.int32) # [256] + even = f4[0::2] + odd = f4[1::2] + if PAIR_HI_FIRST: + dst = ((odd << 4) | even).astype(np.uint8) + else: + dst = ((even << 4) | odd).astype(np.uint8) + assert dst.size == DST_ELEMS + + # v2.bin holds the f32 values as bf16 patterns (host reads uint16). + src_buf = src_bf16.astype(" None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, default=Path(".")) + args = parser.parse_args() + generate(args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/kernel.pto b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/kernel.pto new file mode 100644 index 0000000000..e628ade4c9 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/kernel.pto @@ -0,0 +1,57 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// ----------------------------------------------------------------------------- +// case: vmi_new/quant-bf16x2-to-f4x2-contiguous +// family: conversion +// target_ops: pto.vmi.vcvt (bf16x2 -> f4E1M2x2) +// scenarios: bf16 contiguous load -> vinterpret_cast pair -> vcvt -> f4x2 store +// ----------------------------------------------------------------------------- +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_quant_bf16x2_to_f4x2_contiguous_kernel(%src_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c128_i64 = arith.constant 128 : i64 + %c512_i64 = arith.constant 512 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_dst_u8 = pto.castptr %c512_i64 : i64 -> !pto.ptr + %ub_dst_f4 = pto.castptr %c512_i64 : i64 -> !pto.ptr + + // src: 256 bf16 = 512 bytes; dst: 128 f4x2 = 128 bytes. + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c512_i64 + nburst(%c1_i64, %c512_i64, %c512_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %dst_gm, %ub_dst_u8, %c0_i64, %c128_i64 + nburst(%c1_i64, %c128_i64, %c128_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + pto.vecscope { + %wide = pto.vmi.vload %ub_src[%c0] : !pto.ptr -> !pto.vmi.vreg<256xbf16> + %pair = pto.vmi.vinterpret_cast %wide + : !pto.vmi.vreg<256xbf16> -> !pto.vmi.vreg<128x!pto.bf16x2> + %packed = pto.vmi.vcvt %pair {rounding = "R"} + : !pto.vmi.vreg<128x!pto.bf16x2> -> !pto.vmi.vreg<128x!pto.f4E1M2x2> + pto.vmi.vstore %packed, %ub_dst_f4[%c0] + : !pto.vmi.vreg<128x!pto.f4E1M2x2>, !pto.ptr + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.mte_ub_gm %ub_dst_u8, %dst_gm, %c128_i64 + nburst(%c1_i64, %c128_i64, %c128_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/launch.cpp b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/launch.cpp new file mode 100644 index 0000000000..008fbe5d71 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/launch.cpp @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef __VEC_SCOPE__ +#define __VEC_SCOPE__ +#endif +#if defined(__CCE_AICORE__) && defined(__NPU_ARCH__) && (__NPU_ARCH__ == 2201) +typedef struct { unsigned char v; } hifloat8_t; +typedef struct { unsigned char v; } float8_e4m3_t; +typedef struct { unsigned char v; } float8_e5m2_t; +typedef struct { unsigned char v; } float8_e8m0_t; +typedef struct { unsigned char v; } float4_e1m2x2_t; +typedef struct { unsigned char v; } float4_e2m1x2_t; +#endif +#include +#if !defined(__CCE_AICORE__) && !defined(TMRGSORT_HPP) +struct MrgSortExecutedNumList { + uint16_t mrgSortList0; + uint16_t mrgSortList1; + uint16_t mrgSortList2; + uint16_t mrgSortList3; +}; +#endif +#ifndef __CPU_SIM +#include "acl/acl.h" +#endif + +extern "C" __global__ [aicore] void +vmi_quant_bf16x2_to_f4x2_contiguous_kernel(__gm__ uint16_t *src, __gm__ uint8_t *dst); + +void LaunchVmi_quant_bf16x2_to_f4x2_contiguous_kernel(uint16_t *src, uint8_t *dst, + void *stream) { + vmi_quant_bf16x2_to_f4x2_contiguous_kernel<<<1, nullptr, stream>>>( + (__gm__ uint16_t *)src, (__gm__ uint8_t *)dst); +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/main.cpp b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/main.cpp new file mode 100644 index 0000000000..c6ef1e9171 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/main.cpp @@ -0,0 +1,79 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include + +using namespace PtoTestCommon; + +#define ACL_CHECK(expr) \ + do { \ + const aclError _ret = (expr); \ + if (_ret != ACL_SUCCESS) { \ + std::fprintf(stderr, "[ERROR] %s failed: %d (%s:%d)\n", #expr, \ + (int)_ret, __FILE__, __LINE__); \ + rc = 1; \ + goto cleanup; \ + } \ + } while (0) + +void LaunchVmi_quant_bf16x2_to_f4x2_contiguous_kernel(uint16_t *src, uint8_t *dst, + void *stream); + +int main() { + constexpr size_t kSrcElems = 256; // bf16 + constexpr size_t kDstElems = 128; // f4x2 (1 byte per pair) + size_t srcBytes = kSrcElems * sizeof(uint16_t); + size_t dstBytes = kDstElems * sizeof(uint8_t); + uint16_t *srcHost = nullptr; + uint16_t *srcDevice = nullptr; + uint8_t *dstHost = nullptr; + uint8_t *dstDevice = nullptr; + int rc = 0; + bool aclInited = false; + bool deviceSet = false; + int deviceId = 0; + aclrtStream stream = nullptr; + + ACL_CHECK(aclInit(nullptr)); + aclInited = true; + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) + deviceId = std::atoi(envDevice); + ACL_CHECK(aclrtSetDevice(deviceId)); + deviceSet = true; + ACL_CHECK(aclrtCreateStream(&stream)); + ACL_CHECK(aclrtMallocHost((void **)(&srcHost), srcBytes)); + ACL_CHECK(aclrtMallocHost((void **)(&dstHost), dstBytes)); + ACL_CHECK(aclrtMalloc((void **)&srcDevice, srcBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + ACL_CHECK(aclrtMalloc((void **)&dstDevice, dstBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + + ReadFile("./v1.bin", srcBytes, srcHost, srcBytes); + ReadFile("./v2.bin", dstBytes, dstHost, dstBytes); + ACL_CHECK(aclrtMemcpy(srcDevice, srcBytes, srcHost, srcBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + ACL_CHECK(aclrtMemcpy(dstDevice, dstBytes, dstHost, dstBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + LaunchVmi_quant_bf16x2_to_f4x2_contiguous_kernel(srcDevice, dstDevice, stream); + ACL_CHECK(aclrtSynchronizeStream(stream)); + ACL_CHECK(aclrtMemcpy(dstHost, dstBytes, dstDevice, dstBytes, ACL_MEMCPY_DEVICE_TO_HOST)); + WriteFile("./v2.bin", dstHost, dstBytes); + +cleanup: + aclrtFree(srcDevice); + aclrtFree(dstDevice); + aclrtFreeHost(srcHost); + aclrtFreeHost(dstHost); + if (stream) + aclrtDestroyStream(stream); + if (deviceSet) + aclrtResetDevice(deviceId); + if (aclInited) + aclFinalize(); + return rc; +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/ptoas.flags b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/ptoas.flags new file mode 100644 index 0000000000..5d9dc67120 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-contiguous/ptoas.flags @@ -0,0 +1 @@ +--pto-arch a5 --pto-backend=vpto diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/compare.py b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/compare.py new file mode 100644 index 0000000000..d95d3fec3c --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/compare.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import sys + +import numpy as np + + +def main() -> None: + golden = np.fromfile("golden_v2.bin", dtype=np.uint8) + output = np.fromfile("v2.bin", dtype=np.uint8) + if golden.shape != output.shape: + print(f"[ERROR] size mismatch golden={golden.size} output={output.size}") + sys.exit(2) + if not np.array_equal(golden, output): + diff = np.nonzero(golden != output)[0] + idx = int(diff[0]) + print( + f"[ERROR] compare failed idx={idx} " + f"golden=0x{int(golden[idx]):02x} " + f"output=0x{int(output[idx]):02x}" + ) + sys.exit(2) + print("[INFO] compare passed") + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/golden.py b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/golden.py new file mode 100644 index 0000000000..f2d0e17889 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/golden.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# Reference for bf16x2 -> f4E1M2x2 (full 256-lane, bf16x2 direct-load path). +# Same quantization as the contiguous case; the source is loaded as 256 bf16x2 +# (= 512 bf16) and each pair produces one f4E1M2x2 byte. +# Calibration note: the E1M2 table, rounding interpretation, and low/high +# nibble order below are provisional oracles. Confirm them on the first A5 NPU +# run before treating a strict golden mismatch as a compiler regression. + +import argparse +from pathlib import Path + +import numpy as np + +SRC_ELEMS = 512 # bf16 input lanes (256 bf16x2 pairs) +DST_ELEMS = 256 # f4x2 output bytes + +# Magnitude values for E1M2 code 0..7 (code bit3 = sign). IEEE-style table. +E1M2_VALUES = [0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75] +PAIR_HI_FIRST = True # True: byte = (f4(odd) << 4) | f4(even); False: reversed + +# Period 15 is coprime with the 64/128-lane physical partitions. This prevents +# a part/chunk swap in the d4 lowering from being hidden by a short repetition. +VALUES = np.array([ + 0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75, + -0.25, -0.5, -0.75, -1.0, -1.25, -1.5, -1.75, +], dtype=np.float32) + + +def f32_to_bf16_bits(values: np.ndarray) -> np.ndarray: + """Round f32 to nearest bf16 and return its 16-bit pattern.""" + f32 = values.astype(np.float32) + bits = f32.view(np.uint32) + lsb = (bits >> 16) & 1 + rounding = 0x7FFF + lsb + bf16 = (bits + rounding) >> 16 + return bf16.astype(np.uint16) + + +def bf16_to_f32(bits: np.ndarray) -> np.ndarray: + f32 = bits.astype(np.uint32) << 16 + return f32.view(np.float32) + + +def quantize_f4e1m2(values: np.ndarray) -> np.ndarray: + """Nearest f4E1M2 code (0..15) for each f32 value.""" + sign = (values < 0).astype(np.int32) + mag = np.abs(values) + table = np.array(E1M2_VALUES, dtype=np.float32) + idx = np.argmin(np.abs(table[None, :] - mag[:, None]), axis=1) + return ((sign << 3) | idx).astype(np.uint8) + + +def generate(output_dir: Path) -> None: + repeats = (SRC_ELEMS + len(VALUES) - 1) // len(VALUES) + src_f32 = np.tile(VALUES, repeats)[:SRC_ELEMS].astype(np.float32) + src_bf16 = f32_to_bf16_bits(src_f32) + + f4 = quantize_f4e1m2(bf16_to_f32(src_bf16)).astype(np.int32) # [512] + even = f4[0::2] + odd = f4[1::2] + if PAIR_HI_FIRST: + dst = ((odd << 4) | even).astype(np.uint8) + else: + dst = ((even << 4) | odd).astype(np.uint8) + assert dst.size == DST_ELEMS + + src_buf = src_bf16.astype(" None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, default=Path(".")) + args = parser.parse_args() + generate(args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/kernel.pto b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/kernel.pto new file mode 100644 index 0000000000..70e05e2bc1 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/kernel.pto @@ -0,0 +1,58 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// ----------------------------------------------------------------------------- +// case: vmi_new/quant-bf16x2-to-f4x2-full +// family: conversion +// target_ops: pto.vmi.vcvt (bf16x2 -> f4E1M2x2) +// scenarios: bf16x2 direct load (vlds-of-bf16x2) -> vcvt -> f4x2 store; 256-lane +// multichunk path (4 physical source regs -> 4x vcvt{P0-P3} + 3x vor). +// ----------------------------------------------------------------------------- +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_quant_bf16x2_to_f4x2_full_kernel(%src_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c256_i64 = arith.constant 256 : i64 + %c1024_i64 = arith.constant 1024 : i64 + + %ub_src_u8 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_src_bf16x2 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_dst_u8 = pto.castptr %c1024_i64 : i64 -> !pto.ptr + %ub_dst_f4 = pto.castptr %c1024_i64 : i64 -> !pto.ptr + + // src: 256 bf16x2 = 512 bf16 = 1024 bytes; dst: 256 f4x2 = 256 bytes. + pto.mte_gm_ub %src_gm, %ub_src_u8, %c0_i64, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %dst_gm, %ub_dst_u8, %c0_i64, %c256_i64 + nburst(%c1_i64, %c256_i64, %c256_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + pto.vecscope { + %wide = pto.vmi.vload %ub_src_bf16x2[%c0] + : !pto.ptr -> !pto.vmi.vreg<256x!pto.bf16x2> + %packed = pto.vmi.vcvt %wide {rounding = "R"} + : !pto.vmi.vreg<256x!pto.bf16x2> -> !pto.vmi.vreg<256x!pto.f4E1M2x2> + pto.vmi.vstore %packed, %ub_dst_f4[%c0] + : !pto.vmi.vreg<256x!pto.f4E1M2x2>, !pto.ptr + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.mte_ub_gm %ub_dst_u8, %dst_gm, %c256_i64 + nburst(%c1_i64, %c256_i64, %c256_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/launch.cpp b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/launch.cpp new file mode 100644 index 0000000000..d0ddfed92f --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/launch.cpp @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef __VEC_SCOPE__ +#define __VEC_SCOPE__ +#endif +#if defined(__CCE_AICORE__) && defined(__NPU_ARCH__) && (__NPU_ARCH__ == 2201) +typedef struct { unsigned char v; } hifloat8_t; +typedef struct { unsigned char v; } float8_e4m3_t; +typedef struct { unsigned char v; } float8_e5m2_t; +typedef struct { unsigned char v; } float8_e8m0_t; +typedef struct { unsigned char v; } float4_e1m2x2_t; +typedef struct { unsigned char v; } float4_e2m1x2_t; +#endif +#include +#if !defined(__CCE_AICORE__) && !defined(TMRGSORT_HPP) +struct MrgSortExecutedNumList { + uint16_t mrgSortList0; + uint16_t mrgSortList1; + uint16_t mrgSortList2; + uint16_t mrgSortList3; +}; +#endif +#ifndef __CPU_SIM +#include "acl/acl.h" +#endif + +extern "C" __global__ [aicore] void +vmi_quant_bf16x2_to_f4x2_full_kernel(__gm__ uint16_t *src, __gm__ uint8_t *dst); + +void LaunchVmi_quant_bf16x2_to_f4x2_full_kernel(uint16_t *src, uint8_t *dst, + void *stream) { + vmi_quant_bf16x2_to_f4x2_full_kernel<<<1, nullptr, stream>>>( + (__gm__ uint16_t *)src, (__gm__ uint8_t *)dst); +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/main.cpp b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/main.cpp new file mode 100644 index 0000000000..b0cfa8bd89 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/main.cpp @@ -0,0 +1,79 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include + +using namespace PtoTestCommon; + +#define ACL_CHECK(expr) \ + do { \ + const aclError _ret = (expr); \ + if (_ret != ACL_SUCCESS) { \ + std::fprintf(stderr, "[ERROR] %s failed: %d (%s:%d)\n", #expr, \ + (int)_ret, __FILE__, __LINE__); \ + rc = 1; \ + goto cleanup; \ + } \ + } while (0) + +void LaunchVmi_quant_bf16x2_to_f4x2_full_kernel(uint16_t *src, uint8_t *dst, + void *stream); + +int main() { + constexpr size_t kSrcElems = 512; // bf16 (256 bf16x2 pairs) + constexpr size_t kDstElems = 256; // f4x2 (1 byte per pair) + size_t srcBytes = kSrcElems * sizeof(uint16_t); + size_t dstBytes = kDstElems * sizeof(uint8_t); + uint16_t *srcHost = nullptr; + uint16_t *srcDevice = nullptr; + uint8_t *dstHost = nullptr; + uint8_t *dstDevice = nullptr; + int rc = 0; + bool aclInited = false; + bool deviceSet = false; + int deviceId = 0; + aclrtStream stream = nullptr; + + ACL_CHECK(aclInit(nullptr)); + aclInited = true; + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) + deviceId = std::atoi(envDevice); + ACL_CHECK(aclrtSetDevice(deviceId)); + deviceSet = true; + ACL_CHECK(aclrtCreateStream(&stream)); + ACL_CHECK(aclrtMallocHost((void **)(&srcHost), srcBytes)); + ACL_CHECK(aclrtMallocHost((void **)(&dstHost), dstBytes)); + ACL_CHECK(aclrtMalloc((void **)&srcDevice, srcBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + ACL_CHECK(aclrtMalloc((void **)&dstDevice, dstBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + + ReadFile("./v1.bin", srcBytes, srcHost, srcBytes); + ReadFile("./v2.bin", dstBytes, dstHost, dstBytes); + ACL_CHECK(aclrtMemcpy(srcDevice, srcBytes, srcHost, srcBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + ACL_CHECK(aclrtMemcpy(dstDevice, dstBytes, dstHost, dstBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + LaunchVmi_quant_bf16x2_to_f4x2_full_kernel(srcDevice, dstDevice, stream); + ACL_CHECK(aclrtSynchronizeStream(stream)); + ACL_CHECK(aclrtMemcpy(dstHost, dstBytes, dstDevice, dstBytes, ACL_MEMCPY_DEVICE_TO_HOST)); + WriteFile("./v2.bin", dstHost, dstBytes); + +cleanup: + aclrtFree(srcDevice); + aclrtFree(dstDevice); + aclrtFreeHost(srcHost); + aclrtFreeHost(dstHost); + if (stream) + aclrtDestroyStream(stream); + if (deviceSet) + aclrtResetDevice(deviceId); + if (aclInited) + aclFinalize(); + return rc; +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/ptoas.flags b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/ptoas.flags new file mode 100644 index 0000000000..5d9dc67120 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-full/ptoas.flags @@ -0,0 +1 @@ +--pto-arch a5 --pto-backend=vpto diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/compare.py b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/compare.py new file mode 100644 index 0000000000..d95d3fec3c --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/compare.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import sys + +import numpy as np + + +def main() -> None: + golden = np.fromfile("golden_v2.bin", dtype=np.uint8) + output = np.fromfile("v2.bin", dtype=np.uint8) + if golden.shape != output.shape: + print(f"[ERROR] size mismatch golden={golden.size} output={output.size}") + sys.exit(2) + if not np.array_equal(golden, output): + diff = np.nonzero(golden != output)[0] + idx = int(diff[0]) + print( + f"[ERROR] compare failed idx={idx} " + f"golden=0x{int(golden[idx]):02x} " + f"output=0x{int(output[idx]):02x}" + ) + sys.exit(2) + print("[INFO] compare passed") + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/golden.py b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/golden.py new file mode 100644 index 0000000000..76d7b8d201 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/golden.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# Reference for bf16x2 -> f4E1M2x2 with out-of-range inputs. +# This provisional oracle maps to the nearest entry in the assumed finite E1M2 +# table, including the endpoints for out-of-range values. The actual hardware +# encoding behavior, E1M2 table, rounding interpretation, and nibble order must +# be calibrated on the first A5 NPU run. + +import argparse +from pathlib import Path + +import numpy as np + +SRC_ELEMS = 128 # bf16 input lanes +DST_ELEMS = 64 # f4x2 output bytes + +E1M2_VALUES = [0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75] +PAIR_HI_FIRST = True + +# Overflow (|v| > 1.75), exact max, in-range, and underflow inputs. +VALUES = np.array([100.0, -100.0, 5.0, -5.0, 1.9, -1.9, 1.75, -1.75, + 1.0, -1.0, 0.5, -0.5, 0.0, 1e-5, 3.0, -3.0], + dtype=np.float32) + + +def f32_to_bf16_bits(values: np.ndarray) -> np.ndarray: + f32 = values.astype(np.float32) + bits = f32.view(np.uint32) + lsb = (bits >> 16) & 1 + rounding = 0x7FFF + lsb + bf16 = (bits + rounding) >> 16 + return bf16.astype(np.uint16) + + +def bf16_to_f32(bits: np.ndarray) -> np.ndarray: + return (bits.astype(np.uint32) << 16).view(np.float32) + + +def quantize_nearest_table(values: np.ndarray) -> np.ndarray: + """Map each value to the nearest entry in the provisional E1M2 table.""" + sign = (values < 0).astype(np.int32) + mag = np.abs(values) + table = np.array(E1M2_VALUES, dtype=np.float32) + idx = np.argmin(np.abs(table[None, :] - mag[:, None]), axis=1) + return ((sign << 3) | idx).astype(np.uint8) + + +def generate(output_dir: Path) -> None: + repeats = (SRC_ELEMS + len(VALUES) - 1) // len(VALUES) + src_f32 = np.tile(VALUES, repeats)[:SRC_ELEMS].astype(np.float32) + src_bf16 = f32_to_bf16_bits(src_f32) + + f4 = quantize_nearest_table(bf16_to_f32(src_bf16)).astype(np.int32) + even = f4[0::2] + odd = f4[1::2] + if PAIR_HI_FIRST: + dst = ((odd << 4) | even).astype(np.uint8) + else: + dst = ((even << 4) | odd).astype(np.uint8) + assert dst.size == DST_ELEMS + + src_buf = src_bf16.astype(" None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, default=Path(".")) + args = parser.parse_args() + generate(args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/kernel.pto b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/kernel.pto new file mode 100644 index 0000000000..80fc9e896c --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/kernel.pto @@ -0,0 +1,59 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// ----------------------------------------------------------------------------- +// case: vmi_new/quant-bf16x2-to-f4x2-overflow +// family: conversion +// target_ops: pto.vmi.vcvt (bf16x2 -> f4E1M2x2) +// scenarios: out-of-range bf16 inputs (|v| > E1M2 max 1.75) exercise the +// hardware encoding behavior; tiny inputs probe underflow behavior. +// bf16 contiguous load -> vinterpret_cast pair -> vcvt -> store. +// ----------------------------------------------------------------------------- +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_quant_bf16x2_to_f4x2_overflow_kernel(%src_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c64_i64 = arith.constant 64 : i64 + %c256_i64 = arith.constant 256 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_dst_u8 = pto.castptr %c256_i64 : i64 -> !pto.ptr + %ub_dst_f4 = pto.castptr %c256_i64 : i64 -> !pto.ptr + + // src: 128 bf16 = 256 bytes; dst: 64 f4x2 = 64 bytes. + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c256_i64 + nburst(%c1_i64, %c256_i64, %c256_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %dst_gm, %ub_dst_u8, %c0_i64, %c64_i64 + nburst(%c1_i64, %c64_i64, %c64_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + pto.vecscope { + %wide = pto.vmi.vload %ub_src[%c0] : !pto.ptr -> !pto.vmi.vreg<128xbf16> + %pair = pto.vmi.vinterpret_cast %wide + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + %packed = pto.vmi.vcvt %pair {rounding = "R"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + pto.vmi.vstore %packed, %ub_dst_f4[%c0] + : !pto.vmi.vreg<64x!pto.f4E1M2x2>, !pto.ptr + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.mte_ub_gm %ub_dst_u8, %dst_gm, %c64_i64 + nburst(%c1_i64, %c64_i64, %c64_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/launch.cpp b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/launch.cpp new file mode 100644 index 0000000000..e3730a4cbc --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/launch.cpp @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef __VEC_SCOPE__ +#define __VEC_SCOPE__ +#endif +#if defined(__CCE_AICORE__) && defined(__NPU_ARCH__) && (__NPU_ARCH__ == 2201) +typedef struct { unsigned char v; } hifloat8_t; +typedef struct { unsigned char v; } float8_e4m3_t; +typedef struct { unsigned char v; } float8_e5m2_t; +typedef struct { unsigned char v; } float8_e8m0_t; +typedef struct { unsigned char v; } float4_e1m2x2_t; +typedef struct { unsigned char v; } float4_e2m1x2_t; +#endif +#include +#if !defined(__CCE_AICORE__) && !defined(TMRGSORT_HPP) +struct MrgSortExecutedNumList { + uint16_t mrgSortList0; + uint16_t mrgSortList1; + uint16_t mrgSortList2; + uint16_t mrgSortList3; +}; +#endif +#ifndef __CPU_SIM +#include "acl/acl.h" +#endif + +extern "C" __global__ [aicore] void +vmi_quant_bf16x2_to_f4x2_overflow_kernel(__gm__ uint16_t *src, __gm__ uint8_t *dst); + +void LaunchVmi_quant_bf16x2_to_f4x2_overflow_kernel(uint16_t *src, uint8_t *dst, + void *stream) { + vmi_quant_bf16x2_to_f4x2_overflow_kernel<<<1, nullptr, stream>>>( + (__gm__ uint16_t *)src, (__gm__ uint8_t *)dst); +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/main.cpp b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/main.cpp new file mode 100644 index 0000000000..568fcb9808 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/main.cpp @@ -0,0 +1,79 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include + +using namespace PtoTestCommon; + +#define ACL_CHECK(expr) \ + do { \ + const aclError _ret = (expr); \ + if (_ret != ACL_SUCCESS) { \ + std::fprintf(stderr, "[ERROR] %s failed: %d (%s:%d)\n", #expr, \ + (int)_ret, __FILE__, __LINE__); \ + rc = 1; \ + goto cleanup; \ + } \ + } while (0) + +void LaunchVmi_quant_bf16x2_to_f4x2_overflow_kernel(uint16_t *src, uint8_t *dst, + void *stream); + +int main() { + constexpr size_t kSrcElems = 128; // bf16 + constexpr size_t kDstElems = 64; // f4x2 (1 byte per pair) + size_t srcBytes = kSrcElems * sizeof(uint16_t); + size_t dstBytes = kDstElems * sizeof(uint8_t); + uint16_t *srcHost = nullptr; + uint16_t *srcDevice = nullptr; + uint8_t *dstHost = nullptr; + uint8_t *dstDevice = nullptr; + int rc = 0; + bool aclInited = false; + bool deviceSet = false; + int deviceId = 0; + aclrtStream stream = nullptr; + + ACL_CHECK(aclInit(nullptr)); + aclInited = true; + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) + deviceId = std::atoi(envDevice); + ACL_CHECK(aclrtSetDevice(deviceId)); + deviceSet = true; + ACL_CHECK(aclrtCreateStream(&stream)); + ACL_CHECK(aclrtMallocHost((void **)(&srcHost), srcBytes)); + ACL_CHECK(aclrtMallocHost((void **)(&dstHost), dstBytes)); + ACL_CHECK(aclrtMalloc((void **)&srcDevice, srcBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + ACL_CHECK(aclrtMalloc((void **)&dstDevice, dstBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + + ReadFile("./v1.bin", srcBytes, srcHost, srcBytes); + ReadFile("./v2.bin", dstBytes, dstHost, dstBytes); + ACL_CHECK(aclrtMemcpy(srcDevice, srcBytes, srcHost, srcBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + ACL_CHECK(aclrtMemcpy(dstDevice, dstBytes, dstHost, dstBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + LaunchVmi_quant_bf16x2_to_f4x2_overflow_kernel(srcDevice, dstDevice, stream); + ACL_CHECK(aclrtSynchronizeStream(stream)); + ACL_CHECK(aclrtMemcpy(dstHost, dstBytes, dstDevice, dstBytes, ACL_MEMCPY_DEVICE_TO_HOST)); + WriteFile("./v2.bin", dstHost, dstBytes); + +cleanup: + aclrtFree(srcDevice); + aclrtFree(dstDevice); + aclrtFreeHost(srcHost); + aclrtFreeHost(dstHost); + if (stream) + aclrtDestroyStream(stream); + if (deviceSet) + aclrtResetDevice(deviceId); + if (aclInited) + aclFinalize(); + return rc; +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/ptoas.flags b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/ptoas.flags new file mode 100644 index 0000000000..5d9dc67120 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-overflow/ptoas.flags @@ -0,0 +1 @@ +--pto-arch a5 --pto-backend=vpto diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/compare.py b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/compare.py new file mode 100644 index 0000000000..d95d3fec3c --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/compare.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import sys + +import numpy as np + + +def main() -> None: + golden = np.fromfile("golden_v2.bin", dtype=np.uint8) + output = np.fromfile("v2.bin", dtype=np.uint8) + if golden.shape != output.shape: + print(f"[ERROR] size mismatch golden={golden.size} output={output.size}") + sys.exit(2) + if not np.array_equal(golden, output): + diff = np.nonzero(golden != output)[0] + idx = int(diff[0]) + print( + f"[ERROR] compare failed idx={idx} " + f"golden=0x{int(golden[idx]):02x} " + f"output=0x{int(output[idx]):02x}" + ) + sys.exit(2) + print("[INFO] compare passed") + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/golden.py b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/golden.py new file mode 100644 index 0000000000..01ceeba80a --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/golden.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# Reference for bf16x2 -> f4E1M2x2 rounding {R,A,Z} on f4-boundary values. +# Three chunks of 128 bf16 are quantized with RNE / RTA / RTZ respectively +# (the same input set, so the output regions differ only by rounding mode). +# Calibration note: the E1M2 table, R/A/Z interpretation, tie behavior, and +# low/high nibble order below are provisional oracles. Confirm them on the first +# A5 NPU run before treating a strict golden mismatch as a compiler regression. + +import argparse +from pathlib import Path + +import numpy as np + +SRC_ELEMS = 384 # bf16 input lanes (3 chunks of 128) +DST_ELEMS = 192 # f4x2 output bytes (3 regions of 64) + +# Magnitude values for E1M2 code 0..7 (code bit3 = sign). IEEE-style table. +E1M2_VALUES = [0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75] +PAIR_HI_FIRST = True + +# f4-boundary values (exactly halfway between two representable E1M2 values), +# so RNE / RTA / RTZ exercise tie-breaking differences. +VALUES = np.array([0.125, 0.375, 0.625, 1.125, 1.375, 1.625, -0.125, -0.625], + dtype=np.float32) +CHUNK = 128 + + +def f32_to_bf16_bits(values: np.ndarray) -> np.ndarray: + f32 = values.astype(np.float32) + bits = f32.view(np.uint32) + lsb = (bits >> 16) & 1 + rounding = 0x7FFF + lsb + bf16 = (bits + rounding) >> 16 + return bf16.astype(np.uint16) + + +def bf16_to_f32(bits: np.ndarray) -> np.ndarray: + return (bits.astype(np.uint32) << 16).view(np.float32) + + +def quantize_f4e1m2(values: np.ndarray, mode: str) -> np.ndarray: + """Nearest f4E1M2 code; ties (halfway) resolved by rounding mode.""" + table = np.array(E1M2_VALUES, dtype=np.float32) + out = np.zeros(len(values), dtype=np.int32) + for i, v in enumerate(values): + sign = 1 if v < 0 else 0 + mag = abs(float(v)) + d = np.abs(table - mag) + cands = np.where(np.isclose(d, d.min()))[0] + if len(cands) == 1: + idx = int(cands[0]) + else: + lo, hi = int(cands[0]), int(cands[1]) + if mode == "rne": + idx = lo if lo % 2 == 0 else hi + elif mode == "rta": + idx = hi + elif mode == "rtz": + idx = lo + else: + raise ValueError(mode) + out[i] = (sign << 3) | idx + return out.astype(np.uint8) + + +def pack_pairs(f4: np.ndarray) -> np.ndarray: + f4 = f4.astype(np.int32) + even = f4[0::2] + odd = f4[1::2] + if PAIR_HI_FIRST: + return ((odd << 4) | even).astype(np.uint8) + return ((even << 4) | odd).astype(np.uint8) + + +def generate(output_dir: Path) -> None: + repeats = (SRC_ELEMS + len(VALUES) - 1) // len(VALUES) + src_f32 = np.tile(VALUES, repeats)[:SRC_ELEMS].astype(np.float32) + src_bf16 = f32_to_bf16_bits(src_f32) + + modes = ["rne", "rta", "rtz"] + dst = np.concatenate([ + pack_pairs(quantize_f4e1m2(bf16_to_f32(src_bf16)[c * CHUNK:(c + 1) * CHUNK], m)) + for c, m in enumerate(modes) + ]) + assert dst.size == DST_ELEMS + + src_buf = src_bf16.astype(" None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, default=Path(".")) + args = parser.parse_args() + generate(args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/kernel.pto b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/kernel.pto new file mode 100644 index 0000000000..e23e084a36 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/kernel.pto @@ -0,0 +1,85 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// ----------------------------------------------------------------------------- +// case: vmi_new/quant-bf16x2-to-f4x2-rounding +// family: conversion +// target_ops: pto.vmi.vcvt (bf16x2 -> f4E1M2x2) +// scenarios: rounding {R,A,Z} exercise tie-breaking differences on +// f4-boundary input values. Three chunks +// of the same 128-bf16 boundary input are converted with rounding +// R / A / Z and written to dst regions 0/1/2, so the golden can +// verify the hardware honors each rounding mode (RNE/RTA/RTZ). +// H is rejected by the bf16x2 -> f4x2 VMI contract. +// ----------------------------------------------------------------------------- +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_quant_bf16x2_to_f4x2_rounding_kernel(%src_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c64 = arith.constant 64 : index + %c128 = arith.constant 128 : index + %c256 = arith.constant 256 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c192_i64 = arith.constant 192 : i64 + %c768_i64 = arith.constant 768 : i64 + %c1024_i64 = arith.constant 1024 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_dst_u8 = pto.castptr %c1024_i64 : i64 -> !pto.ptr + %ub_dst_f4 = pto.castptr %c1024_i64 : i64 -> !pto.ptr + + // src: 384 bf16 = 768 bytes (3 chunks of 128); dst: 192 bytes (3 x 64 f4x2). + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c768_i64 + nburst(%c1_i64, %c768_i64, %c768_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %dst_gm, %ub_dst_u8, %c0_i64, %c192_i64 + nburst(%c1_i64, %c192_i64, %c192_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + pto.vecscope { + // chunk 0: rounding R (RNE) + %bf0 = pto.vmi.vload %ub_src[%c0] : !pto.ptr -> !pto.vmi.vreg<128xbf16> + %p0 = pto.vmi.vinterpret_cast %bf0 + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + %r0 = pto.vmi.vcvt %p0 {rounding = "R"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + pto.vmi.vstore %r0, %ub_dst_f4[%c0] + : !pto.vmi.vreg<64x!pto.f4E1M2x2>, !pto.ptr + + // chunk 1: rounding A (RTA) + %bf1 = pto.vmi.vload %ub_src[%c128] : !pto.ptr -> !pto.vmi.vreg<128xbf16> + %p1 = pto.vmi.vinterpret_cast %bf1 + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + %r1 = pto.vmi.vcvt %p1 {rounding = "A"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + pto.vmi.vstore %r1, %ub_dst_f4[%c64] + : !pto.vmi.vreg<64x!pto.f4E1M2x2>, !pto.ptr + + // chunk 2: rounding Z (RTZ) + %bf2 = pto.vmi.vload %ub_src[%c256] : !pto.ptr -> !pto.vmi.vreg<128xbf16> + %p2 = pto.vmi.vinterpret_cast %bf2 + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<64x!pto.bf16x2> + %r2 = pto.vmi.vcvt %p2 {rounding = "Z"} + : !pto.vmi.vreg<64x!pto.bf16x2> -> !pto.vmi.vreg<64x!pto.f4E1M2x2> + pto.vmi.vstore %r2, %ub_dst_f4[%c128] + : !pto.vmi.vreg<64x!pto.f4E1M2x2>, !pto.ptr + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.mte_ub_gm %ub_dst_u8, %dst_gm, %c192_i64 + nburst(%c1_i64, %c192_i64, %c192_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/launch.cpp b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/launch.cpp new file mode 100644 index 0000000000..fbdd28d268 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/launch.cpp @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef __VEC_SCOPE__ +#define __VEC_SCOPE__ +#endif +#if defined(__CCE_AICORE__) && defined(__NPU_ARCH__) && (__NPU_ARCH__ == 2201) +typedef struct { unsigned char v; } hifloat8_t; +typedef struct { unsigned char v; } float8_e4m3_t; +typedef struct { unsigned char v; } float8_e5m2_t; +typedef struct { unsigned char v; } float8_e8m0_t; +typedef struct { unsigned char v; } float4_e1m2x2_t; +typedef struct { unsigned char v; } float4_e2m1x2_t; +#endif +#include +#if !defined(__CCE_AICORE__) && !defined(TMRGSORT_HPP) +struct MrgSortExecutedNumList { + uint16_t mrgSortList0; + uint16_t mrgSortList1; + uint16_t mrgSortList2; + uint16_t mrgSortList3; +}; +#endif +#ifndef __CPU_SIM +#include "acl/acl.h" +#endif + +extern "C" __global__ [aicore] void +vmi_quant_bf16x2_to_f4x2_rounding_kernel(__gm__ uint16_t *src, __gm__ uint8_t *dst); + +void LaunchVmi_quant_bf16x2_to_f4x2_rounding_kernel(uint16_t *src, uint8_t *dst, + void *stream) { + vmi_quant_bf16x2_to_f4x2_rounding_kernel<<<1, nullptr, stream>>>( + (__gm__ uint16_t *)src, (__gm__ uint8_t *)dst); +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/main.cpp b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/main.cpp new file mode 100644 index 0000000000..0574c06586 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/main.cpp @@ -0,0 +1,79 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include + +using namespace PtoTestCommon; + +#define ACL_CHECK(expr) \ + do { \ + const aclError _ret = (expr); \ + if (_ret != ACL_SUCCESS) { \ + std::fprintf(stderr, "[ERROR] %s failed: %d (%s:%d)\n", #expr, \ + (int)_ret, __FILE__, __LINE__); \ + rc = 1; \ + goto cleanup; \ + } \ + } while (0) + +void LaunchVmi_quant_bf16x2_to_f4x2_rounding_kernel(uint16_t *src, uint8_t *dst, + void *stream); + +int main() { + constexpr size_t kSrcElems = 384; // bf16 (3 chunks of 128) + constexpr size_t kDstElems = 192; // f4x2 (3 regions of 64) + size_t srcBytes = kSrcElems * sizeof(uint16_t); + size_t dstBytes = kDstElems * sizeof(uint8_t); + uint16_t *srcHost = nullptr; + uint16_t *srcDevice = nullptr; + uint8_t *dstHost = nullptr; + uint8_t *dstDevice = nullptr; + int rc = 0; + bool aclInited = false; + bool deviceSet = false; + int deviceId = 0; + aclrtStream stream = nullptr; + + ACL_CHECK(aclInit(nullptr)); + aclInited = true; + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) + deviceId = std::atoi(envDevice); + ACL_CHECK(aclrtSetDevice(deviceId)); + deviceSet = true; + ACL_CHECK(aclrtCreateStream(&stream)); + ACL_CHECK(aclrtMallocHost((void **)(&srcHost), srcBytes)); + ACL_CHECK(aclrtMallocHost((void **)(&dstHost), dstBytes)); + ACL_CHECK(aclrtMalloc((void **)&srcDevice, srcBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + ACL_CHECK(aclrtMalloc((void **)&dstDevice, dstBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + + ReadFile("./v1.bin", srcBytes, srcHost, srcBytes); + ReadFile("./v2.bin", dstBytes, dstHost, dstBytes); + ACL_CHECK(aclrtMemcpy(srcDevice, srcBytes, srcHost, srcBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + ACL_CHECK(aclrtMemcpy(dstDevice, dstBytes, dstHost, dstBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + LaunchVmi_quant_bf16x2_to_f4x2_rounding_kernel(srcDevice, dstDevice, stream); + ACL_CHECK(aclrtSynchronizeStream(stream)); + ACL_CHECK(aclrtMemcpy(dstHost, dstBytes, dstDevice, dstBytes, ACL_MEMCPY_DEVICE_TO_HOST)); + WriteFile("./v2.bin", dstHost, dstBytes); + +cleanup: + aclrtFree(srcDevice); + aclrtFree(dstDevice); + aclrtFreeHost(srcHost); + aclrtFreeHost(dstHost); + if (stream) + aclrtDestroyStream(stream); + if (deviceSet) + aclrtResetDevice(deviceId); + if (aclInited) + aclFinalize(); + return rc; +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/ptoas.flags b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/ptoas.flags new file mode 100644 index 0000000000..5d9dc67120 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-rounding/ptoas.flags @@ -0,0 +1 @@ +--pto-arch a5 --pto-backend=vpto diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/compare.py b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/compare.py new file mode 100644 index 0000000000..d95d3fec3c --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/compare.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import sys + +import numpy as np + + +def main() -> None: + golden = np.fromfile("golden_v2.bin", dtype=np.uint8) + output = np.fromfile("v2.bin", dtype=np.uint8) + if golden.shape != output.shape: + print(f"[ERROR] size mismatch golden={golden.size} output={output.size}") + sys.exit(2) + if not np.array_equal(golden, output): + diff = np.nonzero(golden != output)[0] + idx = int(diff[0]) + print( + f"[ERROR] compare failed idx={idx} " + f"golden=0x{int(golden[idx]):02x} " + f"output=0x{int(output[idx]):02x}" + ) + sys.exit(2) + print("[INFO] compare passed") + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/golden.py b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/golden.py new file mode 100644 index 0000000000..f4fdb9e1fb --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/golden.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# Reference for bf16x2 -> f4E1M2x2 tail case. +# 1024-bf16 source buffer holds 1000 logical bf16; the kernel converts in 4 +# chunks of 128 pairs and masks off the tail (500 of 512 f4x2 valid). The +# masked-out dst lanes retain the pre-loaded 0xA5 sentinel. +# Calibration note: the E1M2 table, rounding interpretation, and low/high +# nibble order below are provisional oracles pending the first A5 NPU run. + +import argparse +from pathlib import Path + +import numpy as np + +SRC_ELEMS = 1024 # bf16 source buffer +LOGICAL_BF16 = 1000 # logical bf16 +DST_ELEMS = 512 # f4x2 output buffer +LOGICAL_ELEMS = 500 # logical f4x2 +SENTINEL = np.uint8(0xA5) + +E1M2_VALUES = [0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75] +PAIR_HI_FIRST = True + +VALUES = np.array([0.0, 1.0, -1.0, 0.5, -0.5, 1.5, -1.5, 0.25], dtype=np.float32) + + +def f32_to_bf16_bits(values: np.ndarray) -> np.ndarray: + f32 = values.astype(np.float32) + bits = f32.view(np.uint32) + lsb = (bits >> 16) & 1 + rounding = 0x7FFF + lsb + bf16 = (bits + rounding) >> 16 + return bf16.astype(np.uint16) + + +def bf16_to_f32(bits: np.ndarray) -> np.ndarray: + return (bits.astype(np.uint32) << 16).view(np.float32) + + +def quantize_f4e1m2(values: np.ndarray) -> np.ndarray: + sign = (values < 0).astype(np.int32) + mag = np.abs(values) + table = np.array(E1M2_VALUES, dtype=np.float32) + idx = np.argmin(np.abs(table[None, :] - mag[:, None]), axis=1) + return ((sign << 3) | idx).astype(np.uint8) + + +def generate(output_dir: Path) -> None: + repeats = (SRC_ELEMS + len(VALUES) - 1) // len(VALUES) + src_f32 = np.tile(VALUES, repeats)[:SRC_ELEMS].astype(np.float32) + src_bf16 = f32_to_bf16_bits(src_f32) + + # Quantize the logical bf16 (first LOGICAL_BF16) -> LOGICAL_ELEMS f4x2 bytes. + f4 = quantize_f4e1m2(bf16_to_f32(src_bf16[:LOGICAL_BF16])).astype(np.int32) + even = f4[0::2] + odd = f4[1::2] + if PAIR_HI_FIRST: + valid = ((odd << 4) | even).astype(np.uint8) + else: + valid = ((even << 4) | odd).astype(np.uint8) + assert valid.size == LOGICAL_ELEMS + + dst_buf = np.full(DST_ELEMS, SENTINEL, dtype=np.uint8) + golden = np.full(DST_ELEMS, SENTINEL, dtype=np.uint8) + golden[:LOGICAL_ELEMS] = valid + + src_buf = src_bf16.astype(" None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, default=Path(".")) + args = parser.parse_args() + generate(args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/kernel.pto b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/kernel.pto new file mode 100644 index 0000000000..997a288840 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/kernel.pto @@ -0,0 +1,70 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// ----------------------------------------------------------------------------- +// case: vmi_new/quant-bf16x2-to-f4x2-tail +// family: conversion +// target_ops: pto.vmi.vcvt (bf16x2 -> f4E1M2x2) +// scenarios: tail / partial-register handling. 1000 bf16 (= 500 f4x2) processed +// in 4 chunks of 128 pairs with a create_mask covering the remaining +// pairs; the last chunk (116 pairs) exercises predicated store. +// ----------------------------------------------------------------------------- +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_quant_bf16x2_to_f4x2_tail_kernel(%src_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c2 = arith.constant 2 : index + %c128 = arith.constant 128 : index + %c512 = arith.constant 512 : index + %c500 = arith.constant 500 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c512_i64 = arith.constant 512 : i64 + %c2048_i64 = arith.constant 2048 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_dst_u8 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %ub_dst_f4 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + + // src buffer: 1024 bf16 = 2048 bytes (1000 logical); dst: 512 f4x2 = 512 bytes. + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c2048_i64 + nburst(%c1_i64, %c2048_i64, %c2048_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %dst_gm, %ub_dst_u8, %c0_i64, %c512_i64 + nburst(%c1_i64, %c512_i64, %c512_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + pto.vecscope { + %_:1 = scf.for %offset = %c0 to %c512 step %c128 iter_args(%rem = %c500) -> (index) { + %mask = pto.vmi.create_mask %rem : index -> !pto.vmi.mask<128xpred> + %src_off = arith.muli %offset, %c2 : index + %wide = pto.vmi.vload %ub_src[%src_off] : !pto.ptr -> !pto.vmi.vreg<256xbf16> + %pair = pto.vmi.vinterpret_cast %wide + : !pto.vmi.vreg<256xbf16> -> !pto.vmi.vreg<128x!pto.bf16x2> + %packed = pto.vmi.vcvt %pair {rounding = "R"} + : !pto.vmi.vreg<128x!pto.bf16x2> -> !pto.vmi.vreg<128x!pto.f4E1M2x2> + pto.vmi.vstore %packed, %ub_dst_f4[%offset], %mask + : !pto.vmi.vreg<128x!pto.f4E1M2x2>, !pto.ptr, + !pto.vmi.mask<128xpred> + %next = arith.subi %rem, %c128 : index + scf.yield %next : index + } + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.mte_ub_gm %ub_dst_u8, %dst_gm, %c512_i64 + nburst(%c1_i64, %c512_i64, %c512_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/launch.cpp b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/launch.cpp new file mode 100644 index 0000000000..c94e941d70 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/launch.cpp @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef __VEC_SCOPE__ +#define __VEC_SCOPE__ +#endif +#if defined(__CCE_AICORE__) && defined(__NPU_ARCH__) && (__NPU_ARCH__ == 2201) +typedef struct { unsigned char v; } hifloat8_t; +typedef struct { unsigned char v; } float8_e4m3_t; +typedef struct { unsigned char v; } float8_e5m2_t; +typedef struct { unsigned char v; } float8_e8m0_t; +typedef struct { unsigned char v; } float4_e1m2x2_t; +typedef struct { unsigned char v; } float4_e2m1x2_t; +#endif +#include +#if !defined(__CCE_AICORE__) && !defined(TMRGSORT_HPP) +struct MrgSortExecutedNumList { + uint16_t mrgSortList0; + uint16_t mrgSortList1; + uint16_t mrgSortList2; + uint16_t mrgSortList3; +}; +#endif +#ifndef __CPU_SIM +#include "acl/acl.h" +#endif + +extern "C" __global__ [aicore] void +vmi_quant_bf16x2_to_f4x2_tail_kernel(__gm__ uint16_t *src, __gm__ uint8_t *dst); + +void LaunchVmi_quant_bf16x2_to_f4x2_tail_kernel(uint16_t *src, uint8_t *dst, + void *stream) { + vmi_quant_bf16x2_to_f4x2_tail_kernel<<<1, nullptr, stream>>>( + (__gm__ uint16_t *)src, (__gm__ uint8_t *)dst); +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/main.cpp b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/main.cpp new file mode 100644 index 0000000000..ccc7142d8a --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/main.cpp @@ -0,0 +1,79 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include + +using namespace PtoTestCommon; + +#define ACL_CHECK(expr) \ + do { \ + const aclError _ret = (expr); \ + if (_ret != ACL_SUCCESS) { \ + std::fprintf(stderr, "[ERROR] %s failed: %d (%s:%d)\n", #expr, \ + (int)_ret, __FILE__, __LINE__); \ + rc = 1; \ + goto cleanup; \ + } \ + } while (0) + +void LaunchVmi_quant_bf16x2_to_f4x2_tail_kernel(uint16_t *src, uint8_t *dst, + void *stream); + +int main() { + constexpr size_t kSrcElems = 1024; // bf16 source buffer + constexpr size_t kDstElems = 512; // f4x2 output buffer + size_t srcBytes = kSrcElems * sizeof(uint16_t); + size_t dstBytes = kDstElems * sizeof(uint8_t); + uint16_t *srcHost = nullptr; + uint16_t *srcDevice = nullptr; + uint8_t *dstHost = nullptr; + uint8_t *dstDevice = nullptr; + int rc = 0; + bool aclInited = false; + bool deviceSet = false; + int deviceId = 0; + aclrtStream stream = nullptr; + + ACL_CHECK(aclInit(nullptr)); + aclInited = true; + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) + deviceId = std::atoi(envDevice); + ACL_CHECK(aclrtSetDevice(deviceId)); + deviceSet = true; + ACL_CHECK(aclrtCreateStream(&stream)); + ACL_CHECK(aclrtMallocHost((void **)(&srcHost), srcBytes)); + ACL_CHECK(aclrtMallocHost((void **)(&dstHost), dstBytes)); + ACL_CHECK(aclrtMalloc((void **)&srcDevice, srcBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + ACL_CHECK(aclrtMalloc((void **)&dstDevice, dstBytes, ACL_MEM_MALLOC_HUGE_FIRST)); + + ReadFile("./v1.bin", srcBytes, srcHost, srcBytes); + ReadFile("./v2.bin", dstBytes, dstHost, dstBytes); + ACL_CHECK(aclrtMemcpy(srcDevice, srcBytes, srcHost, srcBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + ACL_CHECK(aclrtMemcpy(dstDevice, dstBytes, dstHost, dstBytes, ACL_MEMCPY_HOST_TO_DEVICE)); + LaunchVmi_quant_bf16x2_to_f4x2_tail_kernel(srcDevice, dstDevice, stream); + ACL_CHECK(aclrtSynchronizeStream(stream)); + ACL_CHECK(aclrtMemcpy(dstHost, dstBytes, dstDevice, dstBytes, ACL_MEMCPY_DEVICE_TO_HOST)); + WriteFile("./v2.bin", dstHost, dstBytes); + +cleanup: + aclrtFree(srcDevice); + aclrtFree(dstDevice); + aclrtFreeHost(srcHost); + aclrtFreeHost(dstHost); + if (stream) + aclrtDestroyStream(stream); + if (deviceSet) + aclrtResetDevice(deviceId); + if (aclInited) + aclFinalize(); + return rc; +} diff --git a/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/ptoas.flags b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/ptoas.flags new file mode 100644 index 0000000000..5d9dc67120 --- /dev/null +++ b/test/vpto/cases/vmi_new/quant-bf16x2-to-f4x2-tail/ptoas.flags @@ -0,0 +1 @@ +--pto-arch a5 --pto-backend=vpto From 4dfae91315a877daec85a5ebb1429578744584f9 Mon Sep 17 00:00:00 2001 From: likai00 Date: Tue, 11 Aug 2026 23:30:08 +0800 Subject: [PATCH 102/122] add docs for vcvt --- docs/isa/vmi-isa/06-convert.md | 34 ++++++++++++++++++++++++++++------ 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/docs/isa/vmi-isa/06-convert.md b/docs/isa/vmi-isa/06-convert.md index edd8fcc4f9..17edae3c6d 100644 --- a/docs/isa/vmi-isa/06-convert.md +++ b/docs/isa/vmi-isa/06-convert.md @@ -18,10 +18,10 @@ dispatches to one of seven kinds: 1. **FpWiden** — `fp → fp`, `|dst| > |src|` (e.g. `f16 → f32`, - `bf16 → f32`, `fp8_e4m3 → f16`). + `bf16 → f32`, `fp8_e4m3 → f16`, `f4x2 → bf16x2`). 2. **FpNarrow** — `fp → fp`, `|dst| < |src|` (e.g. `f32 → f16`, - `f32 → bf16`, `f32 → fp8_e4m3`). Same-width `fp → fp` + `f32 → bf16`, `f32 → fp8_e4m3`, `bf16x2 → f4x2`). Same-width `fp → fp` (`|dst| == |src|`, e.g. `bf16 → f16`). 3. **FpToSi** — `fp → signed int`. Supported pairs follow the contract @@ -57,10 +57,10 @@ | Attribute | Values | Valid for | Description | |---|---|---|---| - | `rounding` | `"R"` (nearest-even), `"A"` (away-from-zero), `"H"` (half-up), `"Z"` (toward-zero) | fp narrowing | Rounding mode | - | `saturate` | `"SAT"`, `"NOSAT"` | required for fp-narrow / int-narrow; for fp→si / fp→ui the requirement follows the vcvt contract's `requiresSat` (e.g. `f16→s8` required, `f16→s32` **forbidden** — no overflow possible; same-width `bf16→f16` required) | `SAT` clamps to ±max of the destination type; `NOSAT` performs a direct bit truncation of the result representation. | + | `rounding` | `"R"` (nearest-even), `"A"` (away-from-zero), `"H"` (half-up), `"Z"` (toward-zero); for the `bf16x2→f4x2` contract pair the allowed set is `"R"`,`"A"`,`"F"` (floor), `"C"` (ceil), `"Z"` (toward-zero) — `"H"` is **rejected** | fp narrowing | Rounding mode | + | `saturate` | `"SAT"`, `"NOSAT"` | required for fp-narrow / int-narrow; for fp→si / fp→ui the requirement follows the vcvt contract's `requiresSat` (e.g. `f16→s8` required, `f16→s32` **forbidden** — no overflow possible; same-width `bf16→f16` required); the `bf16x2→f4x2` narrow has `requiresSat=false` — any `saturate` is **forbidden** | `SAT` clamps to ±max of the destination type; `NOSAT` performs a direct bit truncation of the result representation. | -- **datatypes:** Source and destination from `{f32, f16, bf16, fp8_e4m3, fp8_e5m2, i32, i16, i8, ui32, ui16, ui8}` +- **datatypes:** Source and destination from `{f32, f16, bf16, fp8_e4m3, fp8_e5m2, i32, i16, i8, ui32, ui16, ui8}`; packed carrier types `{!pto.bf16x2, !pto.f4E1M2x2, !pto.f4E2M1x2}` for the bf16x2↔f4x2 fp-to-fp pair (see contract `lookupVMIFpToFpContract`). `bf16x2` is **conversion-only** — it may not appear as a compute element type (`vfadd`/`vfmul`/`vcmp`/...). - **lowering to `pto.mi`:** | Conversion | Physical lowering | `#mi` | `dep` | @@ -72,6 +72,8 @@ | fp↔fp same-width (`bf16→f16`) | `K × vcvt` (1:1, no part) | `K` | `1` | | fp→si / fp→ui | per contract pair: same-width 1:1, widen EVEN/ODD, narrow EVEN/ODD+Vor | `K`–`~3K` | `2`–`3` | | int↔int (same width) | `K × vtrc` or `K × vcvt` | `K` | `1` | + | `bf16x2→f4x2` narrow (32→8) | source viewed as raw `bf16` lanes (2 bf16/bf16x2); `vcvt{P0}` 1:1, `rnd` set, **no sat**; reuse prior pairing `vbitcast` when present | `K` | `1` | + | `f4x2→bf16x2` widen (8→32) | `vcvt{P0}` produces `bf16` lanes; result-side `vbitcast` reinterprets them as `bf16x2`; no rnd, no sat | `K` | `1` | - **example:** ```mlir @@ -109,6 +111,18 @@ // f16 → u8 fp-to-ui (unsigned; contract pair, saturate required) %u = pto.vmi.vcvt %x {saturate = "SAT"} : !pto.vmi.vreg<128×f16> -> !pto.vmi.vreg<128×ui8> + + // bf16x2 → f4x2 quantized narrow (rounding required; saturate forbidden; + // bf16x2 arrives via a physical-noop vinterpret_cast pairing of 2 bf16 lanes) + %pair = pto.vmi.vinterpret_cast %b + : !pto.vmi.vreg<128×bf16> -> !pto.vmi.vreg<64×!pto.bf16x2> + %q4 = pto.vmi.vcvt %pair {rounding = "R"} + : !pto.vmi.vreg<64×!pto.bf16x2> -> !pto.vmi.vreg<64×!pto.f4E1M2x2> + + // f4x2 → bf16x2 dequant widen (no rounding, no saturate; bf16x2 is the + // only legal bf16 carrier for f4 dequant; bare f4x2→bf16 is rejected) + %d = pto.vmi.vcvt %f4 + : !pto.vmi.vreg<64×!pto.f4E1M2x2> -> !pto.vmi.vreg<64×!pto.bf16x2> ``` - **notes:** @@ -120,6 +134,14 @@ - Radix-4 (8↔32) is **not** a stacked predicate chain and **not** a UB roundtrip; the 1↔4 lane spread rides data load/store distribution (`UNPK_B*`/`PK4_B32`) or a `vselr` byte-gather. + - `bf16x2` is **conversion-only**: it is rejected as an element type by all + compute verifiers (`vfadd`/`vfmul`/`vfma`/`vcmp`/`vcmps`/...). The only + way to produce/consume `bf16x2` is via `vcvt` against `f4x2`, or a + bit-conserving `vinterpret_cast` against `bf16` lanes. + - The `bf16x2↔f4x2` pair is the only f4 conversion path exposed at VMI. The + physical `pto.vcvt` consumes/produces raw `bf16` lanes; the `bf16x2` + packaging is a `vbitcast` view inserted by lowering (`vinterpret_cast` + from `128×bf16` to `64×!pto.bf16x2` is a physical no-op pairing). --- @@ -152,7 +174,7 @@ | `result` | `!pto.vmi.vreg` | Bit-reinterpreted vector | - **attributes:** *(none)* -- **datatypes:** Any `T_src`, `T_dst` with `L · bitwidth(T_src) == L · bitwidth(T_dst)` +- **datatypes:** Any `T_src`, `T_dst` (including packed PTO types `!pto.bf16x2`, `!pto.f4E1M2x2`, `!pto.f4E2M1x2`) with `L · bitwidth(T_src) == L · bitwidth(T_dst)` - **lowering to `pto.mi`:** ``` K × pto.vbitcast (or no-op if same physical layout) From 63342c6ffa87001587ed859465f9e6aa5d21f693 Mon Sep 17 00:00:00 2001 From: likai00 Date: Wed, 12 Aug 2026 09:10:25 +0800 Subject: [PATCH 103/122] fix test cases for fp4 --- .../kernel.pto | 5 ++ .../dequant-f4x2-to-bf16x2-tail/kernel.pto | 5 ++ .../dequant-f8-to-f32-tail/kernel_bak.pto | 59 ------------------- .../dequant-f8-to-f32-tail/kernel_vmi_v1.pto | 59 ------------------- .../dequant-f8-to-f32-tail/kernel_vmi_v2.pto | 59 ------------------- 5 files changed, 10 insertions(+), 177 deletions(-) delete mode 100644 test/vpto/cases/vmi_new/dequant-f8-to-f32-tail/kernel_bak.pto delete mode 100644 test/vpto/cases/vmi_new/dequant-f8-to-f32-tail/kernel_vmi_v1.pto delete mode 100644 test/vpto/cases/vmi_new/dequant-f8-to-f32-tail/kernel_vmi_v2.pto diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/kernel.pto b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/kernel.pto index 9cddb046bb..0cf4b9c665 100644 --- a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/kernel.pto +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-contiguous/kernel.pto @@ -33,6 +33,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, !pto.ptr, i64, i64, i64, i64, i64 + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.vecscope { %packed = pto.vmi.vload %ub_src[%c0] : !pto.ptr -> !pto.vmi.vreg<256x!pto.f4E1M2x2> %wide = pto.vmi.vcvt %packed @@ -41,6 +44,8 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, !pto.ptr } + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] pto.mte_ub_gm %ub_dst_bf16, %dst_gm, %c1024_i64 nburst(%c1_i64, %c1024_i64, %c1024_i64) : !pto.ptr, !pto.ptr, i64, i64, i64, i64 diff --git a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/kernel.pto b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/kernel.pto index 8c2ce1aaf4..47cdcb02a2 100644 --- a/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/kernel.pto +++ b/test/vpto/cases/vmi_new/dequant-f4x2-to-bf16x2-tail/kernel.pto @@ -36,6 +36,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, !pto.ptr, i64, i64, i64, i64, i64 + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.vecscope { %_:1 = scf.for %offset = %c0 to %c256 step %c64 iter_args(%rem = %c250) -> (index) { %mask = pto.vmi.create_mask %rem : index -> !pto.vmi.mask<64xpred> @@ -50,6 +53,8 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, !pto.ptr, i64, i64, i64, i64 diff --git a/test/vpto/cases/vmi_new/dequant-f8-to-f32-tail/kernel_bak.pto b/test/vpto/cases/vmi_new/dequant-f8-to-f32-tail/kernel_bak.pto deleted file mode 100644 index 123796c3ba..0000000000 --- a/test/vpto/cases/vmi_new/dequant-f8-to-f32-tail/kernel_bak.pto +++ /dev/null @@ -1,59 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { - func.func @vmi_dequant_f8_to_f32_tail_kernel(%src_gm: !pto.ptr, - %dst_gm: !pto.ptr) attributes {pto.kernel} { - %c0 = arith.constant 0 : index - %c256 = arith.constant 256 : index - %c1024 = arith.constant 1024 : index - %c1000 = arith.constant 1000 : index - %c0_i64 = arith.constant 0 : i64 - %c1_i64 = arith.constant 1 : i64 - %c1024_i64 = arith.constant 1024 : i64 - %c4096_i64 = arith.constant 4096 : i64 - %scale = arith.constant 2.000000e+00 : f32 - - %ub_src_u8 = pto.castptr %c0_i64 : i64 -> !pto.ptr - %ub_src_f8 = pto.castptr %c0_i64 : i64 -> !pto.ptr - %ub_dst = pto.castptr %c4096_i64 : i64 -> !pto.ptr - - pto.mte_gm_ub %src_gm, %ub_src_u8, %c0_i64, %c1024_i64 - nburst(%c1_i64, %c1024_i64, %c1024_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 - pto.mte_gm_ub %dst_gm, %ub_dst, %c0_i64, %c4096_i64 - nburst(%c1_i64, %c4096_i64, %c4096_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 - - pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] - pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] - - pto.vecscope { - %_:1 = scf.for %offset = %c0 to %c1024 step %c256 iter_args(%remaining = %c1000) -> (index) { - %mask = pto.vmi.create_mask %remaining : index -> !pto.vmi.mask<256xpred> - %packed = pto.vmi.vload %ub_src_f8[%offset] : !pto.ptr -> !pto.vmi.vreg<256xf8E4M3FN> - %wide = pto.vmi.vcvt %packed : !pto.vmi.vreg<256xf8E4M3FN> -> !pto.vmi.vreg<256xf32> - %scale_vec = pto.vmi.vbrc %scale : f32 -> !pto.vmi.vreg<256xf32> - %out = pto.vmi.vmul %wide, %scale_vec - : !pto.vmi.vreg<256xf32>, !pto.vmi.vreg<256xf32> -> !pto.vmi.vreg<256xf32> - pto.vmi.masked_store %out, %ub_dst[%offset], %mask - : !pto.vmi.vreg<256xf32>, !pto.ptr, !pto.vmi.mask<256xpred> - %next = arith.subi %remaining, %c256 : index - scf.yield %next : index - } - } - - pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] - pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] - pto.mte_ub_gm %ub_dst, %dst_gm, %c4096_i64 - nburst(%c1_i64, %c4096_i64, %c4096_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64 - pto.barrier #pto.pipe - return - } -} diff --git a/test/vpto/cases/vmi_new/dequant-f8-to-f32-tail/kernel_vmi_v1.pto b/test/vpto/cases/vmi_new/dequant-f8-to-f32-tail/kernel_vmi_v1.pto deleted file mode 100644 index bddf6b0f06..0000000000 --- a/test/vpto/cases/vmi_new/dequant-f8-to-f32-tail/kernel_vmi_v1.pto +++ /dev/null @@ -1,59 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { - func.func @vmi_dequant_f8_to_f32_tail_kernel(%src_gm: !pto.ptr, - %dst_gm: !pto.ptr) attributes {pto.kernel} { - %c0 = arith.constant 0 : index - %c256 = arith.constant 256 : index - %c1024 = arith.constant 1024 : index - %c1000 = arith.constant 1000 : index - %c0_i64 = arith.constant 0 : i64 - %c1_i64 = arith.constant 1 : i64 - %c1024_i64 = arith.constant 1024 : i64 - %c4096_i64 = arith.constant 4096 : i64 - %scale = arith.constant 2.000000e+00 : f32 - - %ub_src_u8 = pto.castptr %c0_i64 : i64 -> !pto.ptr - %ub_src_f8 = pto.castptr %c0_i64 : i64 -> !pto.ptr - %ub_dst = pto.castptr %c4096_i64 : i64 -> !pto.ptr - - pto.mte_gm_ub %src_gm, %ub_src_u8, %c0_i64, %c1024_i64 - nburst(%c1_i64, %c1024_i64, %c1024_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 - pto.mte_gm_ub %dst_gm, %ub_dst, %c0_i64, %c4096_i64 - nburst(%c1_i64, %c4096_i64, %c4096_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 - - pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] - pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] - - pto.vecscope { - %_:1 = scf.for %offset = %c0 to %c1024 step %c256 iter_args(%remaining = %c1000) -> (index) { - %mask = pto.vmi.create_mask %remaining : index -> !pto.vmi.mask<256xpred> - %packed = pto.vmi.load %ub_src_f8[%offset] : !pto.ptr -> !pto.vmi.vreg<256xf8E4M3FN> - %wide = pto.vmi.extf %packed : !pto.vmi.vreg<256xf8E4M3FN> -> !pto.vmi.vreg<256xf32> - %scale_vec = pto.vmi.broadcast %scale : f32 -> !pto.vmi.vreg<256xf32> - %out = pto.vmi.mulf %wide, %scale_vec - : !pto.vmi.vreg<256xf32>, !pto.vmi.vreg<256xf32> -> !pto.vmi.vreg<256xf32> - pto.vmi.masked_store %out, %ub_dst[%offset], %mask - : !pto.vmi.vreg<256xf32>, !pto.ptr, !pto.vmi.mask<256xpred> - %next = arith.subi %remaining, %c256 : index - scf.yield %next : index - } - } - - pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] - pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] - pto.mte_ub_gm %ub_dst, %dst_gm, %c4096_i64 - nburst(%c1_i64, %c4096_i64, %c4096_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64 - pto.barrier #pto.pipe - return - } -} diff --git a/test/vpto/cases/vmi_new/dequant-f8-to-f32-tail/kernel_vmi_v2.pto b/test/vpto/cases/vmi_new/dequant-f8-to-f32-tail/kernel_vmi_v2.pto deleted file mode 100644 index 9ac36bfa82..0000000000 --- a/test/vpto/cases/vmi_new/dequant-f8-to-f32-tail/kernel_vmi_v2.pto +++ /dev/null @@ -1,59 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { - func.func @vmi_dequant_f8_to_f32_tail_kernel(%src_gm: !pto.ptr, - %dst_gm: !pto.ptr) attributes {pto.kernel} { - %c0 = arith.constant 0 : index - %c256 = arith.constant 256 : index - %c1024 = arith.constant 1024 : index - %c1000 = arith.constant 1000 : index - %c0_i64 = arith.constant 0 : i64 - %c1_i64 = arith.constant 1 : i64 - %c1024_i64 = arith.constant 1024 : i64 - %c4096_i64 = arith.constant 4096 : i64 - %scale = arith.constant 2.000000e+00 : f32 - %c1000_i32 = arith.constant 1000 : i32 - - %ub_src_u8 = pto.castptr %c0_i64 : i64 -> !pto.ptr - %ub_src_f8 = pto.castptr %c0_i64 : i64 -> !pto.ptr - %ub_dst = pto.castptr %c4096_i64 : i64 -> !pto.ptr - - pto.mte_gm_ub %src_gm, %ub_src_u8, %c0_i64, %c1024_i64 - nburst(%c1_i64, %c1024_i64, %c1024_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 - pto.mte_gm_ub %dst_gm, %ub_dst, %c0_i64, %c4096_i64 - nburst(%c1_i64, %c4096_i64, %c4096_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 - - pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] - pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] - - pto.vecscope { - %_:1 = scf.for %offset = %c0 to %c1024 step %c256 iter_args(%rem = %c1000_i32) -> (i32) { - %mask, %next = pto.vmi.plt %rem : i32 -> !pto.vmi.mask<256xpred>, i32 - %packed = pto.vmi.vload %ub_src_f8[%offset] : !pto.ptr -> !pto.vmi.vreg<256xf8E4M3FN> - %wide = pto.vmi.vcvt %packed : !pto.vmi.vreg<256xf8E4M3FN> -> !pto.vmi.vreg<256xf32> - %scale_vec = pto.vmi.vbrc %scale : f32 -> !pto.vmi.vreg<256xf32> - %out = pto.vmi.vmul %wide, %scale_vec - : !pto.vmi.vreg<256xf32>, !pto.vmi.vreg<256xf32> -> !pto.vmi.vreg<256xf32> - pto.vmi.vstore %out, %ub_dst[%offset], %mask {pmode = "merge"} - : !pto.vmi.vreg<256xf32>, !pto.ptr, !pto.vmi.mask<256xpred> - scf.yield %next : i32 - } - } - - pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] - pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] - pto.mte_ub_gm %ub_dst, %dst_gm, %c4096_i64 - nburst(%c1_i64, %c4096_i64, %c4096_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64 - pto.barrier #pto.pipe - return - } -} From 5f1f532ba1da9de7884d0c839a6966373205df38 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 16 Jul 2026 16:32:18 +0800 Subject: [PATCH 104/122] Add reusable pto.func helpers --- .../03-kernel-entry-and-subkernels.md | 32 +++- ptodsl/docs/user_guide/05-control-flow.md | 16 +- ptodsl/ptodsl/_cache_signature.py | 58 ++++++ ptodsl/ptodsl/_func.py | 80 +++++++++ ptodsl/ptodsl/_kernel_compilation.py | 51 +----- ptodsl/ptodsl/_tracing/runtime.py | 5 + ptodsl/ptodsl/_tracing/session.py | 165 +++++++++++++++++- ptodsl/ptodsl/pto.py | 1 + ptodsl/tests/test_jit_compile.py | 95 ++++++++++ 9 files changed, 446 insertions(+), 57 deletions(-) create mode 100644 ptodsl/ptodsl/_cache_signature.py create mode 100644 ptodsl/ptodsl/_func.py diff --git a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md index a5d5fe0f46..77b697e026 100644 --- a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md +++ b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md @@ -3,10 +3,11 @@ PTODSL provides one kernel decorator (`@pto.jit`) with two roles (`entry=True` / `entry=False`), two compilation backends (`vpto` / `emitc`), and two reusable compute helper decorators (`@pto.tileop` and `@pto.simt`), -plus inline unit-specific context managers. This chapter covers +plus reusable PTODSL helper functions (`@pto.func`) and inline unit-specific +context managers. This chapter covers the `@pto.jit` entry and module contracts, the two programming models, the two -compilation backends, sub-kernel reference, parameter contracts, and boundary -constraints. +compilation backends, helper functions, sub-kernel reference, parameter +contracts, and boundary constraints. ## 3.1 `@pto.jit` — roles, backends, and modes @@ -22,6 +23,7 @@ Decorator overview: mode="explicit" micro-instruction authoring, user-managed staging @pto.tileop Single-core Tile/scalar compute helper with inferred Vector/Cube kind +@pto.func Reusable PTODSL helper, no host/module ABI boundary @pto.simt Explicitly launched SIMT helper with pointer/scalar ABI ``` @@ -59,6 +61,30 @@ manual-address, user-managed staging contract of explicit kernels. (`@pto.tileop` and `@pto.simt`) define sub-kernels that are called from within `@pto.jit` bodies. +`@pto.func` is the lightweight helper boundary for reusable PTODSL code. It +does not create a host-launchable entry, a kernel-module ABI, or a hardware-unit +sub-kernel section. When traced PTODSL code calls a `@pto.func` helper, PTODSL +materializes one helper `func.func` in the caller's active compilation context +and emits `func.call` at call sites. Supported native Python `if` and +`for range(...)` in the helper body use the same AST rewrite path as `@pto.jit` +and named sub-kernels. The helper can return PTODSL runtime values, including +multiple values via a tuple. + +```python +@pto.func +def add_rows(total: pto.i32, rows: pto.i32): + one = pto.const(1, dtype=pto.i32) + for _ in range(rows): + total = total + one + return total + + +@pto.jit(target="a5") +def kernel(rows: pto.i32): + total = add_rows(pto.const(0, dtype=pto.i32), rows) + _ = total +``` + ## 3.2 `entry=True` — host-launchable kernel entry diff --git a/ptodsl/docs/user_guide/05-control-flow.md b/ptodsl/docs/user_guide/05-control-flow.md index 6768a3f909..e05c61bd55 100644 --- a/ptodsl/docs/user_guide/05-control-flow.md +++ b/ptodsl/docs/user_guide/05-control-flow.md @@ -4,7 +4,7 @@ PTODSL uses a **tracing** compilation model. When you call `kernel.compile(...)` This has one critical implication for how you write loops and branches: -- **Python native `for`/`if`** is rewritten to device-side control flow by default in `@pto.jit` bodies and named `@pto.tileop` / `@pto.simt` helpers. A `for i in range(rows)` loop records a device loop, and a runtime `if` records both branches. + - **Python native `for`/`if`** is rewritten to device-side control flow by default in `@pto.jit` bodies, `@pto.func` helpers, and named `@pto.tileop` / `@pto.simt` helpers. A `for i in range(rows)` loop records a device loop, and a runtime `if` records both branches. - **Assign-form Python conditional expressions** such as `x = a if cond else b` are normalized through the same AST rewrite path, so runtime conditions lower to device-side branches before the assignment is merged back into `x`. - **`pto.const_expr` / `pto.static_range`** keep compile-time Python behavior when you want trace-time specialization or unrolling. - **`pto.for_` / `pto.if_`** produce device-side control flow. The loop bound or branch condition can be a runtime value, and the hardware will execute the loop or take the branch dynamically. @@ -255,11 +255,19 @@ This lets you write a single kernel that specializes into different strategies b ## 5.5 Native Python control-flow rewrite -`@pto.jit` rewrites supported native Python control flow before tracing. In the -default mode, plain Python `if` and `for range(...)` in the rewritten scope -become device-side control flow. Use `pto.const_expr(...)` and +`@pto.jit`, `@pto.func`, and named `@pto.cube` / `@pto.simd` / `@pto.simt` +callables rewrite supported native Python control flow before tracing their +bodies. In the default mode, plain Python `if` and `for range(...)` in the +rewritten scope become device-side control flow. Use `pto.const_expr(...)` and `pto.static_range(...)` when you want trace-time behavior. +PTODSL does not recursively rewrite arbitrary undecorated Python callees. If an +external helper should contain runtime native control flow, decorate it with +`@pto.func` or one of the other PTODSL callable decorators. A plain Python +helper is still executed during tracing; static `range(...)` loops in such a +helper unroll at trace time, and runtime loop bounds or branch conditions are +not converted into `scf.for` / `scf.if`. + ### Runtime branches By default, a native Python `if` becomes a device-side conditional: diff --git a/ptodsl/ptodsl/_cache_signature.py b/ptodsl/ptodsl/_cache_signature.py new file mode 100644 index 0000000000..ebfee276ec --- /dev/null +++ b/ptodsl/ptodsl/_cache_signature.py @@ -0,0 +1,58 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""Shared cache-signature helpers for PTODSL tracing frontends.""" + +from __future__ import annotations + + +def closure_cache_signature(fn): + """Return one stable cache signature for the closure state captured by *fn*.""" + try: + import inspect + + closure_vars = inspect.getclosurevars(fn) + except TypeError: + return () + return tuple( + (name, cache_signature_atom(value)) + for name, value in sorted(closure_vars.nonlocals.items()) + ) + + +def cache_signature_atom(value): + """Return one hashable cache-signature atom for arbitrary captured values.""" + cache_signature = getattr(value, "__ptodsl_cache_signature__", None) + if callable(cache_signature): + return ("ptodsl-cache-signature", cache_signature_atom(cache_signature())) + try: + hash(value) + except TypeError: + if isinstance(value, dict): + items = ( + (cache_signature_atom(key), cache_signature_atom(item)) + for key, item in value.items() + ) + return ("dict", tuple(sorted(items, key=repr))) + if isinstance(value, (list, tuple)): + return ( + type(value).__name__, + tuple(cache_signature_atom(item) for item in value), + ) + if isinstance(value, set): + return ( + "set", + tuple(sorted((cache_signature_atom(item) for item in value), key=repr)), + ) + return (type(value).__name__, repr(value)) + return value + + +__all__ = [ + "cache_signature_atom", + "closure_cache_signature", +] diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py new file mode 100644 index 0000000000..62fa9989b4 --- /dev/null +++ b/ptodsl/ptodsl/_func.py @@ -0,0 +1,80 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""``@pto.func`` decorator and reusable callable handle.""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import update_wrapper +import inspect + +from ._ast_rewrite import rewrite_jit_function +from ._cache_signature import closure_cache_signature +from ._tracing import current_runtime + + +@dataclass(frozen=True) +class FuncSpec: + """Declarative metadata for a rewrite-capable reusable PTODSL callable.""" + + symbol_name: str + + +class FuncTemplate: + """Callable decorated PTODSL helper surface.""" + + def __init__(self, spec: FuncSpec, py_fn, *, ast_rewrite: bool = True): + self.spec = spec + self.py_fn = py_fn + self._ast_rewrite = ast_rewrite + self.signature = inspect.signature(py_fn) + update_wrapper(self, py_fn) + + def emit_body(self, *args, **kwargs): + """Emit this helper body into the currently active trace.""" + py_fn = rewrite_jit_function(self.py_fn) if self._ast_rewrite else self.py_fn + return py_fn(*args, **kwargs) + + def __call__(self, *args, **kwargs): + runtime = current_runtime() + if runtime is None: + raise RuntimeError( + "@pto.func helpers may only be called while tracing a compatible PTODSL kernel" + ) + return runtime.dispatch_ptodsl_func_call(self, *args, **kwargs) + + def __ptodsl_cache_signature__(self): + return ( + type(self).__name__, + self.spec.symbol_name, + id(self.py_fn), + self._ast_rewrite, + closure_cache_signature(self.py_fn), + ) + + +def func(fn=None, *, name: str | None = None, ast_rewrite: bool = True): + """Decorate a Python function as a reusable PTODSL callable helper.""" + + def decorator(py_fn): + return FuncTemplate( + FuncSpec(symbol_name=name or py_fn.__name__), + py_fn, + ast_rewrite=ast_rewrite, + ) + + if fn is not None: + return decorator(fn) + return decorator + + +__all__ = [ + "FuncSpec", + "FuncTemplate", + "func", +] diff --git a/ptodsl/ptodsl/_kernel_compilation.py b/ptodsl/ptodsl/_kernel_compilation.py index 4677a02f8f..60a589534f 100644 --- a/ptodsl/ptodsl/_kernel_compilation.py +++ b/ptodsl/ptodsl/_kernel_compilation.py @@ -9,9 +9,8 @@ from __future__ import annotations -import inspect - from ._ast_rewrite import rewrite_jit_function +from ._cache_signature import closure_cache_signature from ._diagnostics import ( jit_source_compile_constexpr_error, kernel_module_compile_error, @@ -112,7 +111,7 @@ def compile(self, **constexpr_bindings): if self._ast_rewrite: kernel_identity = ( kernel_identity, - _closure_cache_signature(self._callback), + closure_cache_signature(self._callback), ) specialization_key = self._kernel_signature.specialization_key( kernel_identity, @@ -176,52 +175,6 @@ def cached_specializations(self): return tuple(self._compiled_cache.values()) -def _closure_cache_signature(fn): - try: - closure_vars = inspect.getclosurevars(fn) - except TypeError: - return () - return tuple( - (name, _cache_signature_atom(value)) - for name, value in sorted(closure_vars.nonlocals.items()) - ) - - -def _cache_signature_atom(value): - cache_signature = getattr(value, "__ptodsl_cache_signature__", None) - if callable(cache_signature): - return ("ptodsl-cache-signature", _cache_signature_atom(cache_signature())) - try: - hash(value) - except TypeError: - if isinstance(value, dict): - items = ( - (_cache_signature_atom(key), _cache_signature_atom(item)) - for key, item in value.items() - ) - return ( - "dict", - tuple(sorted(items, key=repr)), - ) - if isinstance(value, (list, tuple)): - return ( - type(value).__name__, - tuple(_cache_signature_atom(item) for item in value), - ) - if isinstance(value, set): - return ( - "set", - tuple( - sorted( - (_cache_signature_atom(item) for item in value), - key=repr, - ) - ), - ) - return (type(value).__name__, repr(value)) - return value - - __all__ = [ "CompiledKernelHandle", "KernelCompiler", diff --git a/ptodsl/ptodsl/_tracing/runtime.py b/ptodsl/ptodsl/_tracing/runtime.py index d9172d95e2..f1d425e09e 100644 --- a/ptodsl/ptodsl/_tracing/runtime.py +++ b/ptodsl/ptodsl/_tracing/runtime.py @@ -67,6 +67,11 @@ def dispatch_subkernel_call(self, subkernel, *args, **kwargs): return session.lower_helper_subkernel(subkernel, *args, **kwargs) return subkernel.emit_body(*args, **kwargs) + def dispatch_ptodsl_func_call(self, func_template, *args, **kwargs): + """Dispatch one ``@pto.func`` helper call in the active trace.""" + session = require_active_session("@pto.func") + return session.lower_ptodsl_func_call(func_template, *args, **kwargs) + def dispatch_kernel_module_call(self, kernel_handle, *args, **kwargs): """Dispatch one ``@pto.jit(entry=False)`` kernel-module call in the active trace.""" session = require_active_session("@pto.jit(entry=False)") diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index 2298a38748..d3d7873d2d 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -12,6 +12,7 @@ from contextlib import contextmanager from dataclasses import dataclass import hashlib +import inspect from .._diagnostics import ( inline_subkernel_value_escape_error, @@ -19,6 +20,7 @@ physical_section_value_escape_error, subkernel_kernel_kind_mismatch_error, ) +from .._scalar_coercion import coerce_scalar_to_type from .._kernel_signature import RuntimeScalarParameterSpec from .._ops import const from .._surface_values import ( @@ -28,13 +30,14 @@ is_tile_ir_type, unwrap_surface_value, wrap_like_surface_value, + wrap_surface_value, ) from .control_flow import ( build_carry_loop_frame, finish_carry_loop_frame, yield_carry_loop_state, ) -from .._types import _strip_integer_signedness +from .._types import _resolve, _strip_integer_signedness, int1 from .module_builder import create_container_child_module from ptoas.mlir.dialects import arith, func @@ -48,6 +51,7 @@ IntegerType, Operation, StringAttr, + TypeAttr, UnitAttr, ) @@ -60,6 +64,7 @@ class HelperFunctionSpec: arg_types: tuple result_types: tuple = () attributes: tuple[tuple[str, object], ...] = () + identity: tuple = () def cache_key(self) -> tuple: """Return one stable ABI-sensitive cache key for this helper signature.""" @@ -68,6 +73,7 @@ def cache_key(self) -> tuple: tuple(str(arg_type) for arg_type in self.arg_types), tuple(str(result_type) for result_type in self.result_types), tuple((attr_name, str(attr_value)) for attr_name, attr_value in self.attributes), + self.identity, ) def specialized_symbol_name(self) -> str: @@ -685,6 +691,81 @@ def lower_helper_subkernel(self, subkernel, *args, **kwargs): func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_templates]) return None + def lower_ptodsl_func_call(self, func_template, *args, **kwargs): + """Lower one ``@pto.func`` helper call in the active trace.""" + bound = func_template.signature.bind(*args, **kwargs) + bound.apply_defaults() + ordered_arg_values = [] + normalized_values = {} + for name, param in func_template.signature.parameters.items(): + value = self._normalize_ptodsl_func_argument( + name, + param, + bound.arguments[name], + ) + normalized_values[name] = value + ordered_arg_values.append(value) + if param.kind not in { + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + }: + raise TypeError("@pto.func helpers do not support var-positional or var-keyword parameters yet") + + arg_templates = tuple(ordered_arg_values) + owner_symbol_name = self.current_function_owner_symbol_name + helper_spec = HelperFunctionSpec( + symbol_name=func_template.spec.symbol_name, + arg_types=tuple(unwrap_surface_value(arg).type for arg in arg_templates), + attributes=(("pto.ptodsl.callable_kind", StringAttr.get("func")),), + identity=func_template.__ptodsl_cache_signature__(), + ) + helper_fn, created = self.get_or_create_helper_function( + helper_spec, + owner_symbol_name=owner_symbol_name, + ) + + if created: + entry_block = helper_fn.add_entry_block() + entry_args = tuple(entry_block.arguments) + wrapped_args = [] + wrapped_kwargs = {} + entry_arg_index = 0 + for name, param in func_template.signature.parameters.items(): + entry_arg = entry_args[entry_arg_index] + entry_arg_index += 1 + wrapped_value = wrap_like_surface_value(normalized_values[name], entry_arg) + if param.kind in {inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD}: + wrapped_args.append(wrapped_value) + elif param.kind == inspect.Parameter.KEYWORD_ONLY: + wrapped_kwargs[name] = wrapped_value + else: + raise TypeError("@pto.func helpers do not support var-positional or var-keyword parameters yet") + with ( + self.enter_function(helper_fn, owner_symbol_name=owner_symbol_name), + self.suspend_subkernel_scope(), + InsertionPoint(entry_block), + ): + result = func_template.emit_body(*wrapped_args, **wrapped_kwargs) + return_values = self._normalize_ptodsl_func_return_values( + result, + func_template=func_template, + ) + result_types = tuple(value.type for value in return_values) + helper_fn.attributes["function_type"] = TypeAttr.get( + func.FunctionType.get( + list(helper_spec.arg_types), + list(result_types), + ) + ) + if return_values: + func.ReturnOp([unwrap_surface_value(value) for value in return_values]) + else: + func.ReturnOp([]) + + call_op = func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_templates]) + return self._wrap_ptodsl_func_call_results(call_op.results) + def begin_carry_loop(self, start, stop, step, state_items): """Materialize one authored ``pto.for_(...).carry(...)`` loop body.""" frame = build_carry_loop_frame(start, stop, step, state_items) @@ -905,6 +986,88 @@ def lookup_helper(self, symbol_name: str): return helper return None + def _normalize_ptodsl_func_argument(self, name: str, param, value): + raw_value = unwrap_surface_value(value) + if hasattr(raw_value, "type"): + return raw_value + if param.annotation is not inspect.Parameter.empty: + try: + target_type = _resolve(param.annotation) + except Exception: + target_type = param.annotation + try: + return coerce_scalar_to_type( + raw_value, + target_type, + context=f"@pto.func parameter {name!r}", + ) + except TypeError: + pass + if isinstance(raw_value, bool): + return const(int(raw_value), dtype=int1) + if isinstance(raw_value, int): + return const(raw_value) + raise TypeError( + f"@pto.func parameter {name!r} expects a traced runtime value or a supported literal, " + f"got {raw_value!r}" + ) + + def _normalize_ptodsl_func_return_values(self, result, *, func_template): + if result is None: + return () + if isinstance(result, tuple): + values = result + elif isinstance(result, list): + values = tuple(result) + else: + values = (result,) + + normalized = [] + return_annotation = func_template.signature.return_annotation + target_type = None + if return_annotation is not inspect.Signature.empty: + try: + target_type = _resolve(return_annotation) + except Exception: + target_type = None + + for index, value in enumerate(values): + raw_value = unwrap_surface_value(value) + if hasattr(raw_value, "type"): + normalized.append(raw_value) + continue + if target_type is not None: + try: + normalized.append( + coerce_scalar_to_type( + raw_value, + target_type, + context=f"@pto.func return value {index}", + ) + ) + continue + except TypeError: + pass + if isinstance(raw_value, bool): + normalized.append(const(int(raw_value), dtype=int1)) + continue + if isinstance(raw_value, int): + normalized.append(const(raw_value)) + continue + raise TypeError( + f"@pto.func return value {index} must be a traced runtime value or supported literal, " + f"got {raw_value!r}" + ) + return tuple(normalized) + + def _wrap_ptodsl_func_call_results(self, results): + if not results: + return None + wrapped = tuple(wrap_surface_value(result) for result in results) + if len(wrapped) == 1: + return wrapped[0] + return wrapped + def _attach_ptodsl_logical_name_attr(self, func_op, logical_name: str) -> None: """Mark one ABI-specialized PTODSL symbol with its authored logical name.""" func_op.attributes["pto.ptodsl.logical_name"] = StringAttr.get(logical_name) diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index c3d8a1a4c2..227aaa6558 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -169,6 +169,7 @@ # ── Decorator ───────────────────────────────────────────────────────────────── from ._jit import jit, KernelHandle, merge_jit_modules # noqa: F401 +from ._func import func # noqa: F401 from ._subkernels import cube, simd, simt, tileop # noqa: F401 from ._pipe_namespace import pipe # noqa: F401 diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index ab75b1700c..36a5c6d7dc 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1590,6 +1590,46 @@ def helper(limit, enabled): _ = value +@pto.func +def func_runtime_for_return_helper(limit: pto.i32, initial: pto.i32): + one = pto.const(1, dtype=pto.i32) + total = initial + for _ in range(limit): + total = total + one + return total + + +@pto.func +def func_runtime_if_return_helper(lhs: pto.i32, rhs: pto.i32): + if lhs > rhs: + total = lhs + rhs + else: + total = rhs + lhs + return total + + +@pto.func +def func_multi_return_helper(value: pto.i32): + one = pto.const(1, dtype=pto.i32) + return value, value + one + + +@pto.func +def func_void_helper(): + pto.pipe_barrier(pto.Pipe.ALL) + + +@pto.jit(target="a5") +def ptodsl_func_call_probe(rows: pto.i32): + init = pto.const(0, dtype=pto.i32) + total = func_runtime_for_return_helper(rows, init) + merged = func_runtime_if_return_helper(total, init) + first, second = func_multi_return_helper(merged) + _ = first + second + func_void_helper() + func_void_helper() + + @pto.jit(target="a5") def ast_nested_helper_freevar_if_merge_probe(): lhs = pto.const(4, dtype=pto.i32) @@ -1784,6 +1824,29 @@ def sourceless_subkernel_entry_probe(*, TRACE_TOKEN: pto.const_expr = 0): sourceless_subkernel_entry_probe = make_sourceless_subkernel_entry() +def make_sourceless_ptodsl_func_probe(): + namespace = {"pto": pto} + exec( + """ +@pto.func +def sourceless_ptodsl_func_helper(): + if True: + pto.pipe_barrier(pto.Pipe.ALL) +""", + namespace, + ) + helper = namespace["sourceless_ptodsl_func_helper"] + + @pto.jit(target="a5") + def sourceless_ptodsl_func_probe(*, TRACE_TOKEN: pto.const_expr = 0): + helper() + + return sourceless_ptodsl_func_probe + + +sourceless_ptodsl_func_probe = make_sourceless_ptodsl_func_probe() + + def make_entry_closure_kernel_module_probe(): @pto.jit(target="a5", entry=False) def closure_helper(): @@ -6187,6 +6250,28 @@ def _enter_inline_simt_with_resource_attr(): "rewritten nested helpers should preserve loop-carried and branch live-out values", ) + ptodsl_func_call_text = ptodsl_func_call_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify(ptodsl_func_call_text, "@pto.func helper call specialization") + expect( + re.search(r"func\.func @func_runtime_for_return_helper__ptodsl_[0-9a-f]+\(.*\) -> i32", ptodsl_func_call_text) + is not None, + "@pto.func helpers that return one runtime value should materialize a typed helper result", + ) + expect( + re.search(r"func\.func @func_multi_return_helper__ptodsl_[0-9a-f]+\(.*\) -> \(i32, i32\)", ptodsl_func_call_text) + is not None, + "@pto.func helpers should support multiple returned runtime values", + ) + expect( + ptodsl_func_call_text.count("scf.for") >= 1 and ptodsl_func_call_text.count("scf.if") >= 1, + "@pto.func helper bodies should use native control-flow AST rewrite", + ) + expect( + len(re.findall(r"func\.func @func_void_helper__ptodsl_[0-9a-f]+", ptodsl_func_call_text)) == 1 + and len(re.findall(r"call @func_void_helper__ptodsl_[0-9a-f]+", ptodsl_func_call_text)) == 2, + "repeated @pto.func calls should reuse one materialized helper artifact", + ) + ast_nested_helper_freevar_if_merge_text = ast_nested_helper_freevar_if_merge_probe.compile().mlir_text() expect_parse_roundtrip_and_verify( ast_nested_helper_freevar_if_merge_text, @@ -6310,6 +6395,16 @@ def _enter_inline_simt_with_resource_attr(): "source-less subkernels should fall back to original trace-time Python execution", ) + sourceless_ptodsl_func_text = sourceless_ptodsl_func_probe.compile(TRACE_TOKEN=1).mlir_text() + expect_parse_roundtrip_and_verify( + sourceless_ptodsl_func_text, + "source-less @pto.func AST rewrite fallback specialization", + ) + expect( + sourceless_ptodsl_func_text.count("pto.barrier ") == 1, + "source-less @pto.func helpers should fall back to original trace-time Python execution", + ) + ast_python_bool_guard_enabled_text = ast_python_bool_guard_probe.compile().mlir_text() expect_parse_roundtrip_and_verify( ast_python_bool_guard_enabled_text, From cf1475b5c9fd7df4c8dbad68f3cc45fb4b340f40 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 16 Jul 2026 17:20:00 +0800 Subject: [PATCH 105/122] Add pto.func chain probe artifacts --- tmp/ptodsl_func_chain_probe.mlir | 78 ++++++++++++++++++++++++++++ tmp/ptodsl_func_chain_probe.py | 89 ++++++++++++++++++++++++++++++++ 2 files changed, 167 insertions(+) create mode 100644 tmp/ptodsl_func_chain_probe.mlir create mode 100644 tmp/ptodsl_func_chain_probe.py diff --git a/tmp/ptodsl_func_chain_probe.mlir b/tmp/ptodsl_func_chain_probe.mlir new file mode 100644 index 0000000000..1ff5babfa9 --- /dev/null +++ b/tmp/ptodsl_func_chain_probe.mlir @@ -0,0 +1,78 @@ +// Generated from tmp/ptodsl_func_chain_probe.py while developing issue #946. +// The __ptodsl_ suffixes are specialization hashes and may differ across runs. +// +// Observed behavior: +// 1. Plain undecorated helper: the internal `if True + range(2)` executes during tracing. +// The caller contains two expanded `pto.barrier ` ops and no helper func. +// 2. @pto.func dynamic loop: `for _ in range(limit)` is AST-rewritten and the helper body contains `scf.for`. +// 3. @pto.func dynamic if: `if lhs > rhs` is AST-rewritten and the helper body contains `scf.if`. +// 4. @pto.func(ast_rewrite=False): static `if True + range(2)` does not generate scf. +// It trace-time expands inside the helper body into two `arith.addi` ops. +// 5. Chained calls: `multi_return_helper -> chain_mid -> dyn_loop_helper/dyn_if_helper/no_rewrite_static_helper`. +// 6. Multiple returns: `multi_return_helper` returns `(i32, i32)`. +// 7. Reuse: `dyn_if_helper` is defined once and called twice. +// +// Counts: +// scf.for: 1 +// scf.if: 1 +// func.func @dyn_loop_helper__ptodsl_: 1 +// func.func @dyn_if_helper__ptodsl_: 1 +// func.func @no_rewrite_static_helper__ptodsl_: 1 +// func.func @chain_mid__ptodsl_: 1 +// func.func @multi_return_helper__ptodsl_: 1 +// call @dyn_if_helper__ptodsl_: 2 +// pto.barrier : 2 + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "vpto", pto.kernel_kind = #pto.kernel_kind, pto.target_arch = "a5"} { + func.func @func_chain_probe(%arg0: i32) attributes {pto.entry} { + %c0_i32 = arith.constant 0 : i32 + %0:2 = call @multi_return_helper__ptodsl_406ab2a269(%arg0, %c0_i32) : (i32, i32) -> (i32, i32) + %1 = call @dyn_if_helper__ptodsl_82f7abb427(%0#0, %0#1) : (i32, i32) -> i32 + pto.barrier + pto.barrier + return + } + func.func @multi_return_helper__ptodsl_406ab2a269(%arg0: i32, %arg1: i32) -> (i32, i32) attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "multi_return_helper"} { + %0 = call @chain_mid__ptodsl_be1027dd8d(%arg0, %arg1) : (i32, i32) -> i32 + %c1_i32 = arith.constant 1 : i32 + %1 = arith.addi %0, %c1_i32 : i32 + return %0, %1 : i32, i32 + } + func.func @chain_mid__ptodsl_be1027dd8d(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "chain_mid"} { + %0 = call @dyn_loop_helper__ptodsl_795afe8017(%arg0, %arg1) : (i32, i32) -> i32 + %1 = call @dyn_if_helper__ptodsl_82f7abb427(%0, %arg1) : (i32, i32) -> i32 + %2 = call @no_rewrite_static_helper__ptodsl_eccc8a4ce9(%1) : (i32) -> i32 + return %2 : i32 + } + func.func @dyn_loop_helper__ptodsl_795afe8017(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_loop_helper"} { + %c1_i32 = arith.constant 1 : i32 + %c0 = arith.constant 0 : index + %0 = arith.index_cast %arg0 : i32 to index + %c1 = arith.constant 1 : index + %1 = scf.for %arg2 = %c0 to %0 step %c1 iter_args(%arg3 = %arg1) -> (i32) { + %2 = arith.addi %arg3, %c1_i32 : i32 + scf.yield %2 : i32 + } + return %1 : i32 + } + func.func @dyn_if_helper__ptodsl_82f7abb427(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_if_helper"} { + %0 = arith.cmpi sgt, %arg0, %arg1 : i32 + %1 = scf.if %0 -> (i32) { + %2 = arith.subi %arg0, %arg1 : i32 + scf.yield %2 : i32 + } else { + %2 = arith.subi %arg1, %arg0 : i32 + scf.yield %2 : i32 + } + return %1 : i32 + } + func.func @no_rewrite_static_helper__ptodsl_eccc8a4ce9(%arg0: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "no_rewrite_static_helper"} { + %c1_i32 = arith.constant 1 : i32 + %0 = arith.addi %arg0, %c1_i32 : i32 + %c1_i32_0 = arith.constant 1 : i32 + %1 = arith.addi %0, %c1_i32_0 : i32 + return %1 : i32 + } + } +} diff --git a/tmp/ptodsl_func_chain_probe.py b/tmp/ptodsl_func_chain_probe.py new file mode 100644 index 0000000000..74aec3803e --- /dev/null +++ b/tmp/ptodsl_func_chain_probe.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""Probe for mixed ``@pto.func`` AST rewrite and trace-time expansion behavior.""" + +from ptodsl import pto + + +def plain_trace_helper(): + if True: + for _ in range(2): + pto.pipe_barrier(pto.Pipe.ALL) + + +@pto.func +def dyn_loop_helper(limit: pto.i32, value: pto.i32): + one = pto.const(1, dtype=pto.i32) + total = value + for _ in range(limit): + total = total + one + return total + + +@pto.func +def dyn_if_helper(lhs: pto.i32, rhs: pto.i32): + if lhs > rhs: + chosen = lhs - rhs + else: + chosen = rhs - lhs + return chosen + + +@pto.func(ast_rewrite=False) +def no_rewrite_static_helper(value: pto.i32): + total = value + if True: + for _ in range(2): + total = total + pto.const(1, dtype=pto.i32) + return total + + +@pto.func +def chain_mid(limit: pto.i32, seed: pto.i32): + looped = dyn_loop_helper(limit, seed) + branched = dyn_if_helper(looped, seed) + static_expanded = no_rewrite_static_helper(branched) + return static_expanded + + +@pto.func +def multi_return_helper(limit: pto.i32, seed: pto.i32): + value = chain_mid(limit, seed) + return value, value + pto.const(1, dtype=pto.i32) + + +@pto.jit(target="a5") +def func_chain_probe(limit: pto.i32): + zero = pto.const(0, dtype=pto.i32) + first, second = multi_return_helper(limit, zero) + merged = dyn_if_helper(first, second) + _ = merged + plain_trace_helper() + + +def main(): + text = func_chain_probe.compile().mlir_text() + print(text) + print("\n=== COUNTS ===") + for needle in [ + "scf.for", + "scf.if", + "func.func @dyn_loop_helper__ptodsl_", + "func.func @dyn_if_helper__ptodsl_", + "func.func @no_rewrite_static_helper__ptodsl_", + "func.func @chain_mid__ptodsl_", + "func.func @multi_return_helper__ptodsl_", + "call @dyn_if_helper__ptodsl_", + "pto.barrier ", + ]: + print(f"{needle}: {text.count(needle)}") + + +if __name__ == "__main__": + main() From d152d3f32af01eebfb5924b83b700e3b0387e62b Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 16 Jul 2026 18:51:06 +0800 Subject: [PATCH 106/122] Refine pto.func docs and probe header --- .../03-kernel-entry-and-subkernels.md | 28 +++++++++---------- tmp/ptodsl_func_chain_probe.mlir | 7 +++++ 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md index 77b697e026..eb93cb5357 100644 --- a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md +++ b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md @@ -2,12 +2,12 @@ PTODSL provides one kernel decorator (`@pto.jit`) with two roles (`entry=True` / `entry=False`), two compilation backends (`vpto` / `emitc`), -and two reusable compute helper decorators (`@pto.tileop` and `@pto.simt`), -plus reusable PTODSL helper functions (`@pto.func`) and inline unit-specific -context managers. This chapter covers -the `@pto.jit` entry and module contracts, the two programming models, the two -compilation backends, helper functions, sub-kernel reference, parameter -contracts, and boundary constraints. +two reusable compute helper decorators (`@pto.tileop` and `@pto.simt`), and +reusable PTODSL helper functions (`@pto.func`) plus inline unit-specific +context managers. This chapter covers the `@pto.jit` entry and module +contracts, the two programming models, the two compilation backends, helper +functions, sub-kernel reference, parameter contracts, and boundary +constraints. ## 3.1 `@pto.jit` — roles, backends, and modes @@ -23,7 +23,7 @@ Decorator overview: mode="explicit" micro-instruction authoring, user-managed staging @pto.tileop Single-core Tile/scalar compute helper with inferred Vector/Cube kind -@pto.func Reusable PTODSL helper, no host/module ABI boundary +@pto.func Reusable PTODSL helper with AST-rewritten control flow @pto.simt Explicitly launched SIMT helper with pointer/scalar ABI ``` @@ -61,14 +61,12 @@ manual-address, user-managed staging contract of explicit kernels. (`@pto.tileop` and `@pto.simt`) define sub-kernels that are called from within `@pto.jit` bodies. -`@pto.func` is the lightweight helper boundary for reusable PTODSL code. It -does not create a host-launchable entry, a kernel-module ABI, or a hardware-unit -sub-kernel section. When traced PTODSL code calls a `@pto.func` helper, PTODSL -materializes one helper `func.func` in the caller's active compilation context -and emits `func.call` at call sites. Supported native Python `if` and -`for range(...)` in the helper body use the same AST rewrite path as `@pto.jit` -and named sub-kernels. The helper can return PTODSL runtime values, including -multiple values via a tuple. +`@pto.func` defines reusable PTODSL helper functions. Use it when a helper +contains PTODSL operations or native Python `if` / `for range(...)` that should +compile as device-side control flow. By default, supported native control flow +in a `@pto.func` body is AST-rewritten just like in `@pto.jit` and named +sub-kernels. The helper can return PTODSL runtime values, including multiple +values via a tuple. ```python @pto.func diff --git a/tmp/ptodsl_func_chain_probe.mlir b/tmp/ptodsl_func_chain_probe.mlir index 1ff5babfa9..0915e839b3 100644 --- a/tmp/ptodsl_func_chain_probe.mlir +++ b/tmp/ptodsl_func_chain_probe.mlir @@ -1,3 +1,10 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. // Generated from tmp/ptodsl_func_chain_probe.py while developing issue #946. // The __ptodsl_ suffixes are specialization hashes and may differ across runs. // From e51ad781ad6219022fbabe5e4fe010c9dacbe887 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 16 Jul 2026 18:54:16 +0800 Subject: [PATCH 107/122] Fix pto.func license header --- ptodsl/ptodsl/_func.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py index 62fa9989b4..5e07a911a5 100644 --- a/ptodsl/ptodsl/_func.py +++ b/ptodsl/ptodsl/_func.py @@ -2,7 +2,7 @@ # This program is free software, you can redistribute it and/or modify it under the terms and conditions of # CANN Open Software License Agreement Version 2.0 (the "License"). # Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. """``@pto.func`` decorator and reusable callable handle.""" From 624fd87504ef4229f535d86d8dbac55f4672a05e Mon Sep 17 00:00:00 2001 From: andodo Date: Fri, 17 Jul 2026 09:30:02 +0800 Subject: [PATCH 108/122] Require explicit pto.func return types --- .../03-kernel-entry-and-subkernels.md | 6 +- ptodsl/docs/user_guide/05-control-flow.md | 9 ++- ptodsl/ptodsl/_func.py | 20 ++++- ptodsl/ptodsl/_tracing/session.py | 74 ++++++++++--------- ptodsl/tests/test_jit_compile.py | 15 ++-- tmp/ptodsl_func_chain_probe.mlir | 22 +++--- tmp/ptodsl_func_chain_probe.py | 10 +-- 7 files changed, 90 insertions(+), 66 deletions(-) diff --git a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md index eb93cb5357..fb46df33f5 100644 --- a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md +++ b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md @@ -66,10 +66,12 @@ contains PTODSL operations or native Python `if` / `for range(...)` that should compile as device-side control flow. By default, supported native control flow in a `@pto.func` body is AST-rewritten just like in `@pto.jit` and named sub-kernels. The helper can return PTODSL runtime values, including multiple -values via a tuple. +values via a tuple. Every `@pto.func` helper must declare its return type with +`returns=...` or a Python return annotation; use `returns=None` or `-> None` for +helpers that do not return values. ```python -@pto.func +@pto.func(returns=pto.i32) def add_rows(total: pto.i32, rows: pto.i32): one = pto.const(1, dtype=pto.i32) for _ in range(rows): diff --git a/ptodsl/docs/user_guide/05-control-flow.md b/ptodsl/docs/user_guide/05-control-flow.md index e05c61bd55..3706cfb6fd 100644 --- a/ptodsl/docs/user_guide/05-control-flow.md +++ b/ptodsl/docs/user_guide/05-control-flow.md @@ -263,10 +263,11 @@ rewritten scope become device-side control flow. Use `pto.const_expr(...)` and PTODSL does not recursively rewrite arbitrary undecorated Python callees. If an external helper should contain runtime native control flow, decorate it with -`@pto.func` or one of the other PTODSL callable decorators. A plain Python -helper is still executed during tracing; static `range(...)` loops in such a -helper unroll at trace time, and runtime loop bounds or branch conditions are -not converted into `scf.for` / `scf.if`. +`@pto.func` or one of the other PTODSL callable decorators. `@pto.func` helpers +must explicitly declare their return type with `returns=...` or a Python return +annotation. A plain Python helper is still executed during tracing; static +`range(...)` loops in such a helper unroll at trace time, and runtime loop +bounds or branch conditions are not converted into `scf.for` / `scf.if`. ### Runtime branches diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py index 5e07a911a5..e834efe735 100644 --- a/ptodsl/ptodsl/_func.py +++ b/ptodsl/ptodsl/_func.py @@ -14,9 +14,11 @@ import inspect from ._ast_rewrite import rewrite_jit_function -from ._cache_signature import closure_cache_signature +from ._cache_signature import cache_signature_atom, closure_cache_signature from ._tracing import current_runtime +_RETURNS_UNSET = object() + @dataclass(frozen=True) class FuncSpec: @@ -28,11 +30,21 @@ class FuncSpec: class FuncTemplate: """Callable decorated PTODSL helper surface.""" - def __init__(self, spec: FuncSpec, py_fn, *, ast_rewrite: bool = True): + def __init__(self, spec: FuncSpec, py_fn, *, ast_rewrite: bool = True, returns=_RETURNS_UNSET): self.spec = spec self.py_fn = py_fn self._ast_rewrite = ast_rewrite self.signature = inspect.signature(py_fn) + if returns is not _RETURNS_UNSET: + self.declared_returns = returns + elif self.signature.return_annotation is not inspect.Signature.empty: + self.declared_returns = self.signature.return_annotation + else: + raise TypeError( + "@pto.func helpers must explicitly declare return types with " + "@pto.func(returns=...) or a Python return annotation; use " + "returns=None or -> None for helpers that do not return values" + ) update_wrapper(self, py_fn) def emit_body(self, *args, **kwargs): @@ -54,11 +66,12 @@ def __ptodsl_cache_signature__(self): self.spec.symbol_name, id(self.py_fn), self._ast_rewrite, + cache_signature_atom(self.declared_returns), closure_cache_signature(self.py_fn), ) -def func(fn=None, *, name: str | None = None, ast_rewrite: bool = True): +def func(fn=None, *, name: str | None = None, ast_rewrite: bool = True, returns=_RETURNS_UNSET): """Decorate a Python function as a reusable PTODSL callable helper.""" def decorator(py_fn): @@ -66,6 +79,7 @@ def decorator(py_fn): FuncSpec(symbol_name=name or py_fn.__name__), py_fn, ast_rewrite=ast_rewrite, + returns=returns, ) if fn is not None: diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index d3d7873d2d..c0d3d4d035 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -51,7 +51,6 @@ IntegerType, Operation, StringAttr, - TypeAttr, UnitAttr, ) @@ -717,6 +716,7 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): helper_spec = HelperFunctionSpec( symbol_name=func_template.spec.symbol_name, arg_types=tuple(unwrap_surface_value(arg).type for arg in arg_templates), + result_types=self._declared_ptodsl_func_result_types(func_template), attributes=(("pto.ptodsl.callable_kind", StringAttr.get("func")),), identity=func_template.__ptodsl_cache_signature__(), ) @@ -750,13 +750,7 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): return_values = self._normalize_ptodsl_func_return_values( result, func_template=func_template, - ) - result_types = tuple(value.type for value in return_values) - helper_fn.attributes["function_type"] = TypeAttr.get( - func.FunctionType.get( - list(helper_spec.arg_types), - list(result_types), - ) + result_types=helper_spec.result_types, ) if return_values: func.ReturnOp([unwrap_surface_value(value) for value in return_values]) @@ -1012,8 +1006,21 @@ def _normalize_ptodsl_func_argument(self, name: str, param, value): f"got {raw_value!r}" ) - def _normalize_ptodsl_func_return_values(self, result, *, func_template): + def _declared_ptodsl_func_result_types(self, func_template): + declared_returns = func_template.declared_returns + if declared_returns is None or declared_returns is type(None): + return () + if isinstance(declared_returns, (tuple, list)): + return tuple(_resolve(return_type) for return_type in declared_returns) + return (_resolve(declared_returns),) + + def _normalize_ptodsl_func_return_values(self, result, *, func_template, result_types): if result is None: + if result_types: + raise TypeError( + f"@pto.func {func_template.spec.symbol_name!r} must return " + f"{len(result_types)} value(s) matching its declared return type" + ) return () if isinstance(result, tuple): values = result @@ -1022,41 +1029,36 @@ def _normalize_ptodsl_func_return_values(self, result, *, func_template): else: values = (result,) - normalized = [] - return_annotation = func_template.signature.return_annotation - target_type = None - if return_annotation is not inspect.Signature.empty: - try: - target_type = _resolve(return_annotation) - except Exception: - target_type = None + if len(values) != len(result_types): + raise TypeError( + f"@pto.func {func_template.spec.symbol_name!r} returned {len(values)} value(s), " + f"but its declared return type expects {len(result_types)}" + ) - for index, value in enumerate(values): + normalized = [] + for index, (value, target_type) in enumerate(zip(values, result_types)): raw_value = unwrap_surface_value(value) if hasattr(raw_value, "type"): + if str(raw_value.type) != str(target_type): + raise TypeError( + f"@pto.func return value {index} has type {raw_value.type}, " + f"but the declared return type is {target_type}" + ) normalized.append(raw_value) continue - if target_type is not None: - try: - normalized.append( - coerce_scalar_to_type( - raw_value, - target_type, - context=f"@pto.func return value {index}", - ) + try: + normalized.append( + coerce_scalar_to_type( + raw_value, + target_type, + context=f"@pto.func return value {index}", ) - continue - except TypeError: - pass - if isinstance(raw_value, bool): - normalized.append(const(int(raw_value), dtype=int1)) - continue - if isinstance(raw_value, int): - normalized.append(const(raw_value)) + ) continue + except TypeError: + pass raise TypeError( - f"@pto.func return value {index} must be a traced runtime value or supported literal, " - f"got {raw_value!r}" + f"@pto.func return value {index} must match declared type {target_type}, got {raw_value!r}" ) return tuple(normalized) diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 36a5c6d7dc..b750850dcc 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1590,7 +1590,7 @@ def helper(limit, enabled): _ = value -@pto.func +@pto.func(returns=pto.i32) def func_runtime_for_return_helper(limit: pto.i32, initial: pto.i32): one = pto.const(1, dtype=pto.i32) total = initial @@ -1599,7 +1599,7 @@ def func_runtime_for_return_helper(limit: pto.i32, initial: pto.i32): return total -@pto.func +@pto.func(returns=pto.i32) def func_runtime_if_return_helper(lhs: pto.i32, rhs: pto.i32): if lhs > rhs: total = lhs + rhs @@ -1608,13 +1608,13 @@ def func_runtime_if_return_helper(lhs: pto.i32, rhs: pto.i32): return total -@pto.func +@pto.func(returns=(pto.i32, pto.i32)) def func_multi_return_helper(value: pto.i32): one = pto.const(1, dtype=pto.i32) return value, value + one -@pto.func +@pto.func(returns=None) def func_void_helper(): pto.pipe_barrier(pto.Pipe.ALL) @@ -1828,7 +1828,7 @@ def make_sourceless_ptodsl_func_probe(): namespace = {"pto": pto} exec( """ -@pto.func +@pto.func(returns=None) def sourceless_ptodsl_func_helper(): if True: pto.pipe_barrier(pto.Pipe.ALL) @@ -6271,6 +6271,11 @@ def _enter_inline_simt_with_resource_attr(): and len(re.findall(r"call @func_void_helper__ptodsl_[0-9a-f]+", ptodsl_func_call_text)) == 2, "repeated @pto.func calls should reuse one materialized helper artifact", ) + expect_raises( + TypeError, + lambda: pto.func(lambda value: value), + "must explicitly declare return types", + ) ast_nested_helper_freevar_if_merge_text = ast_nested_helper_freevar_if_merge_probe.compile().mlir_text() expect_parse_roundtrip_and_verify( diff --git a/tmp/ptodsl_func_chain_probe.mlir b/tmp/ptodsl_func_chain_probe.mlir index 0915e839b3..4bcc3d85ed 100644 --- a/tmp/ptodsl_func_chain_probe.mlir +++ b/tmp/ptodsl_func_chain_probe.mlir @@ -34,25 +34,25 @@ module attributes {pto.target_arch = "a5"} { module attributes {pto.backend = "vpto", pto.kernel_kind = #pto.kernel_kind, pto.target_arch = "a5"} { func.func @func_chain_probe(%arg0: i32) attributes {pto.entry} { %c0_i32 = arith.constant 0 : i32 - %0:2 = call @multi_return_helper__ptodsl_406ab2a269(%arg0, %c0_i32) : (i32, i32) -> (i32, i32) - %1 = call @dyn_if_helper__ptodsl_82f7abb427(%0#0, %0#1) : (i32, i32) -> i32 + %0:2 = call @multi_return_helper__ptodsl_c1d36bffde(%arg0, %c0_i32) : (i32, i32) -> (i32, i32) + %1 = call @dyn_if_helper__ptodsl_e04f1ff14e(%0#0, %0#1) : (i32, i32) -> i32 pto.barrier pto.barrier return } - func.func @multi_return_helper__ptodsl_406ab2a269(%arg0: i32, %arg1: i32) -> (i32, i32) attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "multi_return_helper"} { - %0 = call @chain_mid__ptodsl_be1027dd8d(%arg0, %arg1) : (i32, i32) -> i32 + func.func @multi_return_helper__ptodsl_c1d36bffde(%arg0: i32, %arg1: i32) -> (i32, i32) attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "multi_return_helper"} { + %0 = call @chain_mid__ptodsl_f1c09ab46a(%arg0, %arg1) : (i32, i32) -> i32 %c1_i32 = arith.constant 1 : i32 %1 = arith.addi %0, %c1_i32 : i32 return %0, %1 : i32, i32 } - func.func @chain_mid__ptodsl_be1027dd8d(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "chain_mid"} { - %0 = call @dyn_loop_helper__ptodsl_795afe8017(%arg0, %arg1) : (i32, i32) -> i32 - %1 = call @dyn_if_helper__ptodsl_82f7abb427(%0, %arg1) : (i32, i32) -> i32 - %2 = call @no_rewrite_static_helper__ptodsl_eccc8a4ce9(%1) : (i32) -> i32 + func.func @chain_mid__ptodsl_f1c09ab46a(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "chain_mid"} { + %0 = call @dyn_loop_helper__ptodsl_41121a518e(%arg0, %arg1) : (i32, i32) -> i32 + %1 = call @dyn_if_helper__ptodsl_e04f1ff14e(%0, %arg1) : (i32, i32) -> i32 + %2 = call @no_rewrite_static_helper__ptodsl_3485521ccf(%1) : (i32) -> i32 return %2 : i32 } - func.func @dyn_loop_helper__ptodsl_795afe8017(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_loop_helper"} { + func.func @dyn_loop_helper__ptodsl_41121a518e(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_loop_helper"} { %c1_i32 = arith.constant 1 : i32 %c0 = arith.constant 0 : index %0 = arith.index_cast %arg0 : i32 to index @@ -63,7 +63,7 @@ module attributes {pto.target_arch = "a5"} { } return %1 : i32 } - func.func @dyn_if_helper__ptodsl_82f7abb427(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_if_helper"} { + func.func @dyn_if_helper__ptodsl_e04f1ff14e(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_if_helper"} { %0 = arith.cmpi sgt, %arg0, %arg1 : i32 %1 = scf.if %0 -> (i32) { %2 = arith.subi %arg0, %arg1 : i32 @@ -74,7 +74,7 @@ module attributes {pto.target_arch = "a5"} { } return %1 : i32 } - func.func @no_rewrite_static_helper__ptodsl_eccc8a4ce9(%arg0: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "no_rewrite_static_helper"} { + func.func @no_rewrite_static_helper__ptodsl_3485521ccf(%arg0: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "no_rewrite_static_helper"} { %c1_i32 = arith.constant 1 : i32 %0 = arith.addi %arg0, %c1_i32 : i32 %c1_i32_0 = arith.constant 1 : i32 diff --git a/tmp/ptodsl_func_chain_probe.py b/tmp/ptodsl_func_chain_probe.py index 74aec3803e..d47af6e29e 100644 --- a/tmp/ptodsl_func_chain_probe.py +++ b/tmp/ptodsl_func_chain_probe.py @@ -17,7 +17,7 @@ def plain_trace_helper(): pto.pipe_barrier(pto.Pipe.ALL) -@pto.func +@pto.func(returns=pto.i32) def dyn_loop_helper(limit: pto.i32, value: pto.i32): one = pto.const(1, dtype=pto.i32) total = value @@ -26,7 +26,7 @@ def dyn_loop_helper(limit: pto.i32, value: pto.i32): return total -@pto.func +@pto.func(returns=pto.i32) def dyn_if_helper(lhs: pto.i32, rhs: pto.i32): if lhs > rhs: chosen = lhs - rhs @@ -35,7 +35,7 @@ def dyn_if_helper(lhs: pto.i32, rhs: pto.i32): return chosen -@pto.func(ast_rewrite=False) +@pto.func(ast_rewrite=False, returns=pto.i32) def no_rewrite_static_helper(value: pto.i32): total = value if True: @@ -44,7 +44,7 @@ def no_rewrite_static_helper(value: pto.i32): return total -@pto.func +@pto.func(returns=pto.i32) def chain_mid(limit: pto.i32, seed: pto.i32): looped = dyn_loop_helper(limit, seed) branched = dyn_if_helper(looped, seed) @@ -52,7 +52,7 @@ def chain_mid(limit: pto.i32, seed: pto.i32): return static_expanded -@pto.func +@pto.func(returns=(pto.i32, pto.i32)) def multi_return_helper(limit: pto.i32, seed: pto.i32): value = chain_mid(limit, seed) return value, value + pto.const(1, dtype=pto.i32) From 0ce0272c78ab46342a8e4c76afeb098625c586c7 Mon Sep 17 00:00:00 2001 From: andodo Date: Tue, 21 Jul 2026 20:43:47 +0800 Subject: [PATCH 109/122] Reject early returns in rewritten control flow --- ptodsl/ptodsl/_ast_rewrite.py | 40 ++++++++++++++++++++++++ ptodsl/tests/test_jit_compile.py | 53 ++++++++++++++++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index 10720b0d78..7560941849 100644 --- a/ptodsl/ptodsl/_ast_rewrite.py +++ b/ptodsl/ptodsl/_ast_rewrite.py @@ -1007,6 +1007,42 @@ def visit_Subscript(self, node): return ast.copy_location(_name(value_name, node.ctx), node) return self.generic_visit(node) +class _ControlFlowExitVisitor(ast.NodeVisitor): + def __init__(self): + self.exit_node = None + + def visit_Return(self, node): + self.exit_node = node + + def visit_Yield(self, node): + self.exit_node = node + + def visit_YieldFrom(self, node): + self.exit_node = node + + def visit_FunctionDef(self, node): + return + + def visit_AsyncFunctionDef(self, node): + return + + def visit_Lambda(self, node): + return + + def visit_ClassDef(self, node): + return + + +def _reject_control_flow_exits(stmts, context: str): + visitor = _ControlFlowExitVisitor() + for stmt in stmts: + visitor.visit(stmt) + if visitor.exit_node is not None: + raise PTODSLAstRewriteError( + f"ast_rewrite=True does not support return/yield inside rewritten {context}; " + "assign values to locals and return after the rewritten control flow" + ) + class _ControlFlowRewriter: def __init__(self, static_env=None, *, section_entry_bindings=None, section_uninitialized_aliases=None): @@ -1151,6 +1187,9 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con ) return [stmt] + _reject_control_flow_exits(stmt.body, "if branches") + _reject_control_flow_exits(stmt.orelse, "if branches") + cond_name = self._fresh("cond") then_info = _name_info(stmt.body) else_info = _name_info(stmt.orelse) @@ -1420,6 +1459,7 @@ def _rewrite_for(self, stmt, *, live_after, live_after_slots=None, allow_loop_co raise PTODSLAstRewriteError("ast_rewrite=True does not support for-else on runtime loops") if not isinstance(stmt.target, ast.Name): raise PTODSLAstRewriteError("ast_rewrite=True runtime for-loops require a simple name target") + _reject_control_flow_exits(stmt.body, "for-loop bodies") if stmt.target.id in live_after: raise PTODSLAstRewriteError( "ast_rewrite=True runtime for-loops cannot expose the loop induction variable outside the loop yet; " diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index b750850dcc..ee79a229c9 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1608,6 +1608,26 @@ def func_runtime_if_return_helper(lhs: pto.i32, rhs: pto.i32): return total +@pto.func(returns=pto.i32) +def func_runtime_if_early_return_helper(lhs: pto.i32, rhs: pto.i32): + if lhs > rhs: + return lhs + return rhs + + +@pto.func(returns=pto.i32) +def func_runtime_for_early_return_helper(limit: pto.i32, initial: pto.i32): + for _ in range(limit): + return initial + return initial + + +@pto.func(returns=None) +def func_runtime_if_yield_helper(lhs: pto.i32, rhs: pto.i32): + if lhs > rhs: + yield lhs + + @pto.func(returns=(pto.i32, pto.i32)) def func_multi_return_helper(value: pto.i32): one = pto.const(1, dtype=pto.i32) @@ -1630,6 +1650,24 @@ def ptodsl_func_call_probe(rows: pto.i32): func_void_helper() +@pto.jit(target="a5") +def ptodsl_func_if_early_return_probe(rows: pto.i32): + init = pto.const(0, dtype=pto.i32) + _ = func_runtime_if_early_return_helper(rows, init) + + +@pto.jit(target="a5") +def ptodsl_func_for_early_return_probe(rows: pto.i32): + init = pto.const(0, dtype=pto.i32) + _ = func_runtime_for_early_return_helper(rows, init) + + +@pto.jit(target="a5") +def ptodsl_func_if_yield_probe(rows: pto.i32): + init = pto.const(0, dtype=pto.i32) + func_runtime_if_yield_helper(rows, init) + + @pto.jit(target="a5") def ast_nested_helper_freevar_if_merge_probe(): lhs = pto.const(4, dtype=pto.i32) @@ -6271,6 +6309,21 @@ def _enter_inline_simt_with_resource_attr(): and len(re.findall(r"call @func_void_helper__ptodsl_[0-9a-f]+", ptodsl_func_call_text)) == 2, "repeated @pto.func calls should reuse one materialized helper artifact", ) + expect_raises( + PTODSLAstRewriteError, + lambda: ptodsl_func_if_early_return_probe.compile().mlir_text(), + "return/yield inside rewritten if branches", + ) + expect_raises( + PTODSLAstRewriteError, + lambda: ptodsl_func_for_early_return_probe.compile().mlir_text(), + "return/yield inside rewritten for-loop bodies", + ) + expect_raises( + PTODSLAstRewriteError, + lambda: ptodsl_func_if_yield_probe.compile().mlir_text(), + "return/yield inside rewritten if branches", + ) expect_raises( TypeError, lambda: pto.func(lambda value: value), From 5a4dde1dee0ec2176152cb4a3bf90e6ca81f292d Mon Sep 17 00:00:00 2001 From: andodo Date: Wed, 22 Jul 2026 09:04:54 +0800 Subject: [PATCH 110/122] Preserve func helper surface templates --- ptodsl/ptodsl/_tracing/session.py | 17 ++++++++++------- ptodsl/tests/test_jit_compile.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 7 deletions(-) diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index c0d3d4d035..ba3671b1fe 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -695,15 +695,16 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): bound = func_template.signature.bind(*args, **kwargs) bound.apply_defaults() ordered_arg_values = [] - normalized_values = {} + arg_templates = [] for name, param in func_template.signature.parameters.items(): + original_value = bound.arguments[name] value = self._normalize_ptodsl_func_argument( name, param, - bound.arguments[name], + original_value, ) - normalized_values[name] = value ordered_arg_values.append(value) + arg_templates.append(original_value) if param.kind not in { inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, @@ -711,11 +712,12 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): }: raise TypeError("@pto.func helpers do not support var-positional or var-keyword parameters yet") - arg_templates = tuple(ordered_arg_values) + arg_values = tuple(ordered_arg_values) + arg_templates = tuple(arg_templates) owner_symbol_name = self.current_function_owner_symbol_name helper_spec = HelperFunctionSpec( symbol_name=func_template.spec.symbol_name, - arg_types=tuple(unwrap_surface_value(arg).type for arg in arg_templates), + arg_types=tuple(unwrap_surface_value(arg).type for arg in arg_values), result_types=self._declared_ptodsl_func_result_types(func_template), attributes=(("pto.ptodsl.callable_kind", StringAttr.get("func")),), identity=func_template.__ptodsl_cache_signature__(), @@ -733,8 +735,9 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): entry_arg_index = 0 for name, param in func_template.signature.parameters.items(): entry_arg = entry_args[entry_arg_index] + arg_template = arg_templates[entry_arg_index] entry_arg_index += 1 - wrapped_value = wrap_like_surface_value(normalized_values[name], entry_arg) + wrapped_value = wrap_like_surface_value(arg_template, entry_arg) if param.kind in {inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD}: wrapped_args.append(wrapped_value) elif param.kind == inspect.Parameter.KEYWORD_ONLY: @@ -757,7 +760,7 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): else: func.ReturnOp([]) - call_op = func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_templates]) + call_op = func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_values]) return self._wrap_ptodsl_func_call_results(call_op.results) def begin_carry_loop(self, start, stop, step, state_items): diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index ee79a229c9..3f3b200ca5 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1639,6 +1639,11 @@ def func_void_helper(): pto.pipe_barrier(pto.Pipe.ALL) +@pto.func(returns=pto.i32) +def func_partition_metadata_helper(part: pto.PartitionTensorView, cols: pto.i32): + return part.sizes[0] + cols + + @pto.jit(target="a5") def ptodsl_func_call_probe(rows: pto.i32): init = pto.const(0, dtype=pto.i32) @@ -1650,6 +1655,16 @@ def ptodsl_func_call_probe(rows: pto.i32): func_void_helper() +@pto.jit(target="a5") +def ptodsl_func_partition_metadata_probe( + A_ptr: pto.ptr(pto.f32, "gm"), + cols: pto.i32, +): + a_view = pto.make_tensor_view(A_ptr, shape=[1, 16], strides=[16, 1]) + part = pto.partition_view(a_view, offsets=[0, 0], sizes=[1, 16]) + _ = func_partition_metadata_helper(part, cols) + + @pto.jit(target="a5") def ptodsl_func_if_early_return_probe(rows: pto.i32): init = pto.const(0, dtype=pto.i32) @@ -6309,6 +6324,19 @@ def _enter_inline_simt_with_resource_attr(): and len(re.findall(r"call @func_void_helper__ptodsl_[0-9a-f]+", ptodsl_func_call_text)) == 2, "repeated @pto.func calls should reuse one materialized helper artifact", ) + ptodsl_func_partition_metadata_text = ptodsl_func_partition_metadata_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + ptodsl_func_partition_metadata_text, + "@pto.func partition metadata specialization", + ) + expect( + re.search(r"func\.func @func_partition_metadata_helper__ptodsl_[0-9a-f]+", ptodsl_func_partition_metadata_text) + is not None + and re.search(r"func\.func @func_partition_metadata_helper__ptodsl_[0-9a-f]+\(.*\) -> i32", ptodsl_func_partition_metadata_text) + is not None + and ptodsl_func_partition_metadata_text.count("pto.partition_view") >= 2, + "@pto.func should preserve partition metadata across the helper boundary", + ) expect_raises( PTODSLAstRewriteError, lambda: ptodsl_func_if_early_return_probe.compile().mlir_text(), From e886a96ff42cf9312aca41d68e31e38187ffc532 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 10:21:42 +0800 Subject: [PATCH 111/122] Support postponed annotations in pto.func --- ptodsl/ptodsl/_func.py | 20 +++++++++ ptodsl/ptodsl/_tracing/session.py | 24 ++++++----- ptodsl/tests/test_jit_compile.py | 72 +++++++++++++++++++++++++++++++ 3 files changed, 105 insertions(+), 11 deletions(-) diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py index e834efe735..8cf2a64602 100644 --- a/ptodsl/ptodsl/_func.py +++ b/ptodsl/ptodsl/_func.py @@ -12,6 +12,7 @@ from dataclasses import dataclass from functools import update_wrapper import inspect +import typing from ._ast_rewrite import rewrite_jit_function from ._cache_signature import cache_signature_atom, closure_cache_signature @@ -35,8 +36,18 @@ def __init__(self, spec: FuncSpec, py_fn, *, ast_rewrite: bool = True, returns=_ self.py_fn = py_fn self._ast_rewrite = ast_rewrite self.signature = inspect.signature(py_fn) + try: + self.type_hints = typing.get_type_hints(py_fn) + except Exception as exc: + if _has_annotations(self.signature): + raise TypeError( + f"failed to resolve @pto.func annotations for {py_fn.__qualname__!r}" + ) from exc + self.type_hints = {} if returns is not _RETURNS_UNSET: self.declared_returns = returns + elif "return" in self.type_hints: + self.declared_returns = self.type_hints["return"] elif self.signature.return_annotation is not inspect.Signature.empty: self.declared_returns = self.signature.return_annotation else: @@ -87,6 +98,15 @@ def decorator(py_fn): return decorator +def _has_annotations(signature: inspect.Signature) -> bool: + if signature.return_annotation is not inspect.Signature.empty: + return True + return any( + param.annotation is not inspect.Parameter.empty + for param in signature.parameters.values() + ) + + __all__ = [ "FuncSpec", "FuncTemplate", diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index ba3671b1fe..148d2b5c22 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -696,21 +696,23 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): bound.apply_defaults() ordered_arg_values = [] arg_templates = [] + type_hints = getattr(func_template, "type_hints", {}) for name, param in func_template.signature.parameters.items(): + if param.kind not in { + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + }: + raise TypeError("@pto.func helpers do not support var-positional or var-keyword parameters yet") original_value = bound.arguments[name] + annotation = type_hints.get(name, param.annotation) value = self._normalize_ptodsl_func_argument( name, - param, + annotation, original_value, ) ordered_arg_values.append(value) arg_templates.append(original_value) - if param.kind not in { - inspect.Parameter.POSITIONAL_ONLY, - inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY, - }: - raise TypeError("@pto.func helpers do not support var-positional or var-keyword parameters yet") arg_values = tuple(ordered_arg_values) arg_templates = tuple(arg_templates) @@ -983,15 +985,15 @@ def lookup_helper(self, symbol_name: str): return helper return None - def _normalize_ptodsl_func_argument(self, name: str, param, value): + def _normalize_ptodsl_func_argument(self, name: str, annotation, value): raw_value = unwrap_surface_value(value) if hasattr(raw_value, "type"): return raw_value - if param.annotation is not inspect.Parameter.empty: + if annotation is not inspect.Parameter.empty: try: - target_type = _resolve(param.annotation) + target_type = _resolve(annotation) except Exception: - target_type = param.annotation + target_type = annotation try: return coerce_scalar_to_type( raw_value, diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 3f3b200ca5..ab8b51209f 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1644,6 +1644,41 @@ def func_partition_metadata_helper(part: pto.PartitionTensorView, cols: pto.i32) return part.sizes[0] + cols +def _make_future_annotations_ptodsl_func_helpers(): + namespace = {"pto": pto} + exec( + """ +from __future__ import annotations + +@pto.func +def func_future_i32_return_helper(x: pto.i32) -> pto.i32: + return x + pto.const(1, dtype=pto.i32) + +@pto.func +def func_future_i64_literal_helper(x: pto.i64) -> pto.i64: + return x + pto.const(1, dtype=pto.i64) + +@pto.func +def func_future_void_helper(x: pto.i32) -> None: + _ = x + pto.pipe_barrier(pto.Pipe.ALL) +""", + namespace, + ) + return ( + namespace["func_future_i32_return_helper"], + namespace["func_future_i64_literal_helper"], + namespace["func_future_void_helper"], + ) + + +( + func_future_i32_return_helper, + func_future_i64_literal_helper, + func_future_void_helper, +) = _make_future_annotations_ptodsl_func_helpers() + + @pto.jit(target="a5") def ptodsl_func_call_probe(rows: pto.i32): init = pto.const(0, dtype=pto.i32) @@ -1665,6 +1700,14 @@ def ptodsl_func_partition_metadata_probe( _ = func_partition_metadata_helper(part, cols) +@pto.jit(target="a5") +def ptodsl_func_future_annotations_probe(rows: pto.i32): + i32_value = func_future_i32_return_helper(rows) + i64_value = func_future_i64_literal_helper(1) + func_future_void_helper(i32_value) + _ = i64_value + + @pto.jit(target="a5") def ptodsl_func_if_early_return_probe(rows: pto.i32): init = pto.const(0, dtype=pto.i32) @@ -6337,6 +6380,35 @@ def _enter_inline_simt_with_resource_attr(): and ptodsl_func_partition_metadata_text.count("pto.partition_view") >= 2, "@pto.func should preserve partition metadata across the helper boundary", ) + ptodsl_func_future_annotations_text = ptodsl_func_future_annotations_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + ptodsl_func_future_annotations_text, + "@pto.func future annotations specialization", + ) + expect( + re.search( + r"func\.func @func_future_i32_return_helper__ptodsl_[0-9a-f]+\(.*i32.*\) -> i32", + ptodsl_func_future_annotations_text, + ) + is not None, + "@pto.func should resolve PEP 563 string return annotations to PTO result types", + ) + expect( + re.search( + r"func\.func @func_future_i64_literal_helper__ptodsl_[0-9a-f]+\(.*i64.*\) -> i64", + ptodsl_func_future_annotations_text, + ) + is not None, + "@pto.func should resolve PEP 563 string parameter annotations before materializing literals", + ) + expect( + re.search( + r"func\.func @func_future_void_helper__ptodsl_[0-9a-f]+\(.*i32.*\) attributes", + ptodsl_func_future_annotations_text, + ) + is not None, + "@pto.func should resolve PEP 563 -> None annotations as void helpers", + ) expect_raises( PTODSLAstRewriteError, lambda: ptodsl_func_if_early_return_probe.compile().mlir_text(), From 34acf75662c2e15881c41d50c57bd40179dc1da4 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 10:32:14 +0800 Subject: [PATCH 112/122] Relax pto.func partition metadata test --- ptodsl/tests/test_jit_compile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index ab8b51209f..4e16bc2068 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -6377,7 +6377,7 @@ def _enter_inline_simt_with_resource_attr(): is not None and re.search(r"func\.func @func_partition_metadata_helper__ptodsl_[0-9a-f]+\(.*\) -> i32", ptodsl_func_partition_metadata_text) is not None - and ptodsl_func_partition_metadata_text.count("pto.partition_view") >= 2, + and re.search(r"%c1_i32 = arith\.constant 1 : i32", ptodsl_func_partition_metadata_text) is not None, "@pto.func should preserve partition metadata across the helper boundary", ) ptodsl_func_future_annotations_text = ptodsl_func_future_annotations_probe.compile().mlir_text() From 35570d39c1c68bd7a18d48ba378ed590c7814db1 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 11:10:31 +0800 Subject: [PATCH 113/122] Stabilize pto.func helper identity --- ptodsl/ptodsl/_cache_signature.py | 18 ++++++++++++++++++ ptodsl/ptodsl/_func.py | 4 ++-- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/ptodsl/ptodsl/_cache_signature.py b/ptodsl/ptodsl/_cache_signature.py index ebfee276ec..c401977786 100644 --- a/ptodsl/ptodsl/_cache_signature.py +++ b/ptodsl/ptodsl/_cache_signature.py @@ -10,6 +10,23 @@ from __future__ import annotations +def function_cache_signature(fn): + """Return one stable cache signature for a Python function body.""" + code = fn.__code__ + return ( + "python-function", + code.co_filename, + code.co_firstlineno, + code.co_qualname, + code.co_argcount, + code.co_posonlyargcount, + code.co_kwonlyargcount, + code.co_names, + code.co_consts, + code.co_code, + ) + + def closure_cache_signature(fn): """Return one stable cache signature for the closure state captured by *fn*.""" try: @@ -55,4 +72,5 @@ def cache_signature_atom(value): __all__ = [ "cache_signature_atom", "closure_cache_signature", + "function_cache_signature", ] diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py index 8cf2a64602..be317087bc 100644 --- a/ptodsl/ptodsl/_func.py +++ b/ptodsl/ptodsl/_func.py @@ -15,7 +15,7 @@ import typing from ._ast_rewrite import rewrite_jit_function -from ._cache_signature import cache_signature_atom, closure_cache_signature +from ._cache_signature import cache_signature_atom, closure_cache_signature, function_cache_signature from ._tracing import current_runtime _RETURNS_UNSET = object() @@ -75,7 +75,7 @@ def __ptodsl_cache_signature__(self): return ( type(self).__name__, self.spec.symbol_name, - id(self.py_fn), + function_cache_signature(self.py_fn), self._ast_rewrite, cache_signature_atom(self.declared_returns), closure_cache_signature(self.py_fn), From 71bf26b0d779c7ddf4ddb464e24d6fb0128eb9c1 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 14:16:12 +0800 Subject: [PATCH 114/122] Preserve helper ABI cache key --- ptodsl/ptodsl/_tracing/session.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index 148d2b5c22..ee903f53dd 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -72,12 +72,17 @@ def cache_key(self) -> tuple: tuple(str(arg_type) for arg_type in self.arg_types), tuple(str(result_type) for result_type in self.result_types), tuple((attr_name, str(attr_value)) for attr_name, attr_value in self.attributes), - self.identity, ) + def specialization_key(self) -> tuple: + """Return one stable semantic cache key for this helper specialization.""" + if not self.identity: + return self.cache_key() + return (*self.cache_key(), self.identity) + def specialized_symbol_name(self) -> str: - """Return one stable symbol name that is unique for this helper ABI.""" - digest = hashlib.sha1(repr(self.cache_key()).encode("utf-8")).hexdigest()[:10] + """Return one stable symbol name that is unique for this helper specialization.""" + digest = hashlib.sha1(repr(self.specialization_key()).encode("utf-8")).hexdigest()[:10] return f"{self.symbol_name}__ptodsl_{digest}" @@ -1089,7 +1094,7 @@ def get_or_create_helper_function(self, spec: HelperFunctionSpec, *, owner_symbo owner_symbol_name = ( self.current_function_owner_symbol_name if owner_symbol_name is None else owner_symbol_name ) - cache_key = (owner_symbol_name, spec.cache_key()) + cache_key = (owner_symbol_name, spec.specialization_key()) helper = self._helpers.get(cache_key) if helper is not None: return helper, False @@ -1112,7 +1117,7 @@ def get_or_create_helper_function(self, spec: HelperFunctionSpec, *, owner_symbo def get_or_create_kernel_module_primary_function(self, spec: HelperFunctionSpec, module_spec): """Look up or create the primary definition for one kernel-module callee.""" - cache_key = spec.cache_key() + cache_key = spec.specialization_key() helper = self._kernel_module_primary_functions.get(cache_key) if helper is not None: return helper, False From 2678ad7100764b3eb55df45f2b26b472099777d3 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 14:19:31 +0800 Subject: [PATCH 115/122] Support older code object metadata --- ptodsl/ptodsl/_cache_signature.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ptodsl/ptodsl/_cache_signature.py b/ptodsl/ptodsl/_cache_signature.py index c401977786..92a111fae0 100644 --- a/ptodsl/ptodsl/_cache_signature.py +++ b/ptodsl/ptodsl/_cache_signature.py @@ -17,7 +17,7 @@ def function_cache_signature(fn): "python-function", code.co_filename, code.co_firstlineno, - code.co_qualname, + getattr(code, "co_qualname", fn.__qualname__), code.co_argcount, code.co_posonlyargcount, code.co_kwonlyargcount, From b7194c999062dfe75ce00d659cb1d9c8e3aad9d0 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 15:48:22 +0800 Subject: [PATCH 116/122] Support const_expr pto.func parameters --- ptodsl/ptodsl/_tracing/session.py | 58 +++++++++++++++++++++++-------- ptodsl/tests/test_jit_compile.py | 44 +++++++++++++++++++++++ 2 files changed, 87 insertions(+), 15 deletions(-) diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index ee903f53dd..e3f5e4b8c8 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -20,9 +20,11 @@ physical_section_value_escape_error, subkernel_kernel_kind_mismatch_error, ) +from .._cache_signature import cache_signature_atom from .._scalar_coercion import coerce_scalar_to_type from .._kernel_signature import RuntimeScalarParameterSpec from .._ops import const +from .._surface_types import const_expr as _const_expr_marker from .._surface_values import ( AddressOffsetValue, AllocatedBufferValue, @@ -699,8 +701,10 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): """Lower one ``@pto.func`` helper call in the active trace.""" bound = func_template.signature.bind(*args, **kwargs) bound.apply_defaults() - ordered_arg_values = [] - arg_templates = [] + runtime_arg_values = [] + runtime_arg_templates = [] + param_bindings = [] + constexpr_bindings = [] type_hints = getattr(func_template, "type_hints", {}) for name, param in func_template.signature.parameters.items(): if param.kind not in { @@ -711,23 +715,35 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): raise TypeError("@pto.func helpers do not support var-positional or var-keyword parameters yet") original_value = bound.arguments[name] annotation = type_hints.get(name, param.annotation) + if annotation is _const_expr_marker: + value = self._normalize_ptodsl_func_constexpr_argument(name, original_value) + param_bindings.append(("constexpr", name, param, value)) + constexpr_bindings.append((name, cache_signature_atom(value))) + continue value = self._normalize_ptodsl_func_argument( name, annotation, original_value, ) - ordered_arg_values.append(value) - arg_templates.append(original_value) - - arg_values = tuple(ordered_arg_values) - arg_templates = tuple(arg_templates) + param_bindings.append(("runtime", name, param, original_value)) + runtime_arg_values.append(value) + runtime_arg_templates.append(original_value) + + runtime_arg_values = tuple(runtime_arg_values) + runtime_arg_templates = tuple(runtime_arg_templates) + identity = func_template.__ptodsl_cache_signature__() + if constexpr_bindings: + identity = ( + identity, + ("constexprs", tuple(constexpr_bindings)), + ) owner_symbol_name = self.current_function_owner_symbol_name helper_spec = HelperFunctionSpec( symbol_name=func_template.spec.symbol_name, - arg_types=tuple(unwrap_surface_value(arg).type for arg in arg_values), + arg_types=tuple(unwrap_surface_value(arg).type for arg in runtime_arg_values), result_types=self._declared_ptodsl_func_result_types(func_template), attributes=(("pto.ptodsl.callable_kind", StringAttr.get("func")),), - identity=func_template.__ptodsl_cache_signature__(), + identity=identity, ) helper_fn, created = self.get_or_create_helper_function( helper_spec, @@ -740,11 +756,14 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): wrapped_args = [] wrapped_kwargs = {} entry_arg_index = 0 - for name, param in func_template.signature.parameters.items(): - entry_arg = entry_args[entry_arg_index] - arg_template = arg_templates[entry_arg_index] - entry_arg_index += 1 - wrapped_value = wrap_like_surface_value(arg_template, entry_arg) + for binding_kind, name, param, value in param_bindings: + if binding_kind == "constexpr": + wrapped_value = value + else: + entry_arg = entry_args[entry_arg_index] + arg_template = runtime_arg_templates[entry_arg_index] + entry_arg_index += 1 + wrapped_value = wrap_like_surface_value(arg_template, entry_arg) if param.kind in {inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD}: wrapped_args.append(wrapped_value) elif param.kind == inspect.Parameter.KEYWORD_ONLY: @@ -767,7 +786,7 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): else: func.ReturnOp([]) - call_op = func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_values]) + call_op = func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in runtime_arg_values]) return self._wrap_ptodsl_func_call_results(call_op.results) def begin_carry_loop(self, start, stop, step, state_items): @@ -1016,6 +1035,15 @@ def _normalize_ptodsl_func_argument(self, name: str, annotation, value): f"got {raw_value!r}" ) + def _normalize_ptodsl_func_constexpr_argument(self, name: str, value): + raw_value = unwrap_surface_value(value) + if hasattr(raw_value, "type"): + raise TypeError( + f"@pto.func const_expr parameter {name!r} expects a compile-time Python value, " + f"got traced runtime value of type {raw_value.type}" + ) + return raw_value + def _declared_ptodsl_func_result_types(self, func_template): declared_returns = func_template.declared_returns if declared_returns is None or declared_returns is type(None): diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 4e16bc2068..224934a40b 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1644,6 +1644,15 @@ def func_partition_metadata_helper(part: pto.PartitionTensorView, cols: pto.i32) return part.sizes[0] + cols +@pto.func(returns=pto.i32) +def func_constexpr_static_helper(value: pto.i32, *, BLOCK: pto.const_expr = 2): + total = value + one = pto.const(1, dtype=pto.i32) + for _ in pto.static_range(BLOCK): + total = total + one + return total + + def _make_future_annotations_ptodsl_func_helpers(): namespace = {"pto": pto} exec( @@ -1708,6 +1717,18 @@ def ptodsl_func_future_annotations_probe(rows: pto.i32): _ = i64_value +@pto.jit(target="a5") +def ptodsl_func_constexpr_probe(rows: pto.i32): + first = func_constexpr_static_helper(rows, BLOCK=2) + second = func_constexpr_static_helper(rows, BLOCK=4) + _ = first + second + + +@pto.jit(target="a5") +def ptodsl_func_constexpr_runtime_value_probe(rows: pto.i32): + _ = func_constexpr_static_helper(rows, BLOCK=rows) + + @pto.jit(target="a5") def ptodsl_func_if_early_return_probe(rows: pto.i32): init = pto.const(0, dtype=pto.i32) @@ -6409,6 +6430,29 @@ def _enter_inline_simt_with_resource_attr(): is not None, "@pto.func should resolve PEP 563 -> None annotations as void helpers", ) + ptodsl_func_constexpr_text = ptodsl_func_constexpr_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + ptodsl_func_constexpr_text, + "@pto.func const_expr specialization", + ) + constexpr_helper_names = re.findall( + r"func\.func @(func_constexpr_static_helper__ptodsl_[0-9a-f]+)\(%arg0: i32\) -> i32", + ptodsl_func_constexpr_text, + ) + expect( + len(set(constexpr_helper_names)) == 2, + "@pto.func const_expr parameters should specialize helpers without entering the runtime ABI", + ) + expect( + ptodsl_func_constexpr_text.count("call @func_constexpr_static_helper__ptodsl_") == 2 + and ptodsl_func_constexpr_text.count("arith.addi") >= 6, + "@pto.func const_expr values should remain available for static_range unrolling", + ) + expect_raises( + TypeError, + lambda: ptodsl_func_constexpr_runtime_value_probe.compile().mlir_text(), + "const_expr parameter 'BLOCK' expects a compile-time Python value", + ) expect_raises( PTODSLAstRewriteError, lambda: ptodsl_func_if_early_return_probe.compile().mlir_text(), From 09756277dfd56cf7f98d67abe251191098f80a74 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 16:25:27 +0800 Subject: [PATCH 117/122] Keep bare returns compatible in AST rewrite --- ptodsl/ptodsl/_ast_rewrite.py | 50 ++++++++++++++++++++++++++++------- ptodsl/ptodsl/_func.py | 6 ++++- 2 files changed, 46 insertions(+), 10 deletions(-) diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index 7560941849..6fe12b6b30 100644 --- a/ptodsl/ptodsl/_ast_rewrite.py +++ b/ptodsl/ptodsl/_ast_rewrite.py @@ -20,7 +20,13 @@ class PTODSLAstRewriteError(SyntaxError): """Raised when AST rewrite sees unsupported Python control flow.""" -def rewrite_jit_function(fn, *, static_bindings=None, rewrite_control_flow=True): +def rewrite_jit_function( + fn, + *, + static_bindings=None, + rewrite_control_flow=True, + reject_bare_returns: bool = False, +): """Return a function with PTODSL lexical sections lowered safely. ``pto.section`` is a physical SSA region, not a Python ``with`` hint. The @@ -28,6 +34,9 @@ def rewrite_jit_function(fn, *, static_bindings=None, rewrite_control_flow=True) the optional control-flow rewrite is disabled. This keeps Python's function-local assignment rules from leaking a section-local SSA value into a sibling physical section. + ``reject_bare_returns`` controls whether a bare ``return`` inside + rewritten control flow is rejected, while value returns are always + rejected because rewritten branches must communicate through locals. """ try: source = inspect.getsource(fn) @@ -60,6 +69,7 @@ def rewrite_jit_function(fn, *, static_bindings=None, rewrite_control_flow=True) static_env, section_entry_bindings=section_rewriter.section_entry_bindings, section_uninitialized_aliases=section_rewriter.section_uninitialized_aliases, + reject_bare_returns=reject_bare_returns, ) function_def.body = rewriter.rewrite_block(function_def.body, live_after=set()) tree = ast.Module(body=[function_def], type_ignores=[]) @@ -1008,11 +1018,13 @@ def visit_Subscript(self, node): return self.generic_visit(node) class _ControlFlowExitVisitor(ast.NodeVisitor): - def __init__(self): + def __init__(self, *, reject_bare_returns: bool): self.exit_node = None + self._reject_bare_returns = reject_bare_returns def visit_Return(self, node): - self.exit_node = node + if self._reject_bare_returns or node.value is not None: + self.exit_node = node def visit_Yield(self, node): self.exit_node = node @@ -1033,8 +1045,8 @@ def visit_ClassDef(self, node): return -def _reject_control_flow_exits(stmts, context: str): - visitor = _ControlFlowExitVisitor() +def _reject_control_flow_exits(stmts, context: str, *, reject_bare_returns: bool): + visitor = _ControlFlowExitVisitor(reject_bare_returns=reject_bare_returns) for stmt in stmts: visitor.visit(stmt) if visitor.exit_node is not None: @@ -1045,11 +1057,19 @@ def _reject_control_flow_exits(stmts, context: str): class _ControlFlowRewriter: - def __init__(self, static_env=None, *, section_entry_bindings=None, section_uninitialized_aliases=None): + def __init__( + self, + static_env=None, + *, + section_entry_bindings=None, + section_uninitialized_aliases=None, + reject_bare_returns: bool = False, + ): self._static_env = dict(static_env or {}) self._section_entry_bindings = dict(section_entry_bindings or {}) self._section_uninitialized_aliases = set(section_uninitialized_aliases or ()) self._counter = 0 + self._reject_bare_returns = reject_bare_returns def _fresh(self, prefix: str) -> str: value = f"__pto_ast_{prefix}_{self._counter}" @@ -1187,8 +1207,16 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con ) return [stmt] - _reject_control_flow_exits(stmt.body, "if branches") - _reject_control_flow_exits(stmt.orelse, "if branches") + _reject_control_flow_exits( + stmt.body, + "if branches", + reject_bare_returns=self._reject_bare_returns, + ) + _reject_control_flow_exits( + stmt.orelse, + "if branches", + reject_bare_returns=self._reject_bare_returns, + ) cond_name = self._fresh("cond") then_info = _name_info(stmt.body) @@ -1459,7 +1487,11 @@ def _rewrite_for(self, stmt, *, live_after, live_after_slots=None, allow_loop_co raise PTODSLAstRewriteError("ast_rewrite=True does not support for-else on runtime loops") if not isinstance(stmt.target, ast.Name): raise PTODSLAstRewriteError("ast_rewrite=True runtime for-loops require a simple name target") - _reject_control_flow_exits(stmt.body, "for-loop bodies") + _reject_control_flow_exits( + stmt.body, + "for-loop bodies", + reject_bare_returns=self._reject_bare_returns, + ) if stmt.target.id in live_after: raise PTODSLAstRewriteError( "ast_rewrite=True runtime for-loops cannot expose the loop induction variable outside the loop yet; " diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py index be317087bc..f79c72ff02 100644 --- a/ptodsl/ptodsl/_func.py +++ b/ptodsl/ptodsl/_func.py @@ -60,7 +60,11 @@ def __init__(self, spec: FuncSpec, py_fn, *, ast_rewrite: bool = True, returns=_ def emit_body(self, *args, **kwargs): """Emit this helper body into the currently active trace.""" - py_fn = rewrite_jit_function(self.py_fn) if self._ast_rewrite else self.py_fn + py_fn = ( + rewrite_jit_function(self.py_fn, reject_bare_returns=True) + if self._ast_rewrite + else self.py_fn + ) return py_fn(*args, **kwargs) def __call__(self, *args, **kwargs): From a791082c5100497460647c7414d11299479f90d6 Mon Sep 17 00:00:00 2001 From: andodo Date: Mon, 10 Aug 2026 16:34:25 +0800 Subject: [PATCH 118/122] Fix pto.func SSA arguments and closure captures --- ptodsl/ptodsl/_func.py | 26 ++++++++++++++++ ptodsl/ptodsl/_tracing/session.py | 13 ++++++-- ptodsl/tests/test_jit_compile.py | 50 +++++++++++++++++++++++++++++++ 3 files changed, 87 insertions(+), 2 deletions(-) diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py index f79c72ff02..55eb931cff 100644 --- a/ptodsl/ptodsl/_func.py +++ b/ptodsl/ptodsl/_func.py @@ -16,6 +16,7 @@ from ._ast_rewrite import rewrite_jit_function from ._cache_signature import cache_signature_atom, closure_cache_signature, function_cache_signature +from ._surface_values import unwrap_surface_value from ._tracing import current_runtime _RETURNS_UNSET = object() @@ -67,12 +68,22 @@ def emit_body(self, *args, **kwargs): ) return py_fn(*args, **kwargs) + def _validate_closure_captures(self): + for name, value in inspect.getclosurevars(self.py_fn).nonlocals.items(): + captured = _find_runtime_value(value) + if captured is not None: + raise TypeError( + f"@pto.func {self.spec.symbol_name!r} captures runtime value {name!r} " + f"of type {captured.type}; pass it as an explicit parameter" + ) + def __call__(self, *args, **kwargs): runtime = current_runtime() if runtime is None: raise RuntimeError( "@pto.func helpers may only be called while tracing a compatible PTODSL kernel" ) + self._validate_closure_captures() return runtime.dispatch_ptodsl_func_call(self, *args, **kwargs) def __ptodsl_cache_signature__(self): @@ -111,6 +122,21 @@ def _has_annotations(signature: inspect.Signature) -> bool: ) +def _find_runtime_value(value): + if isinstance(value, dict): + values = value.items() + elif isinstance(value, (tuple, list, set, frozenset)): + values = value + else: + raw_value = unwrap_surface_value(value) + return raw_value if hasattr(raw_value, "type") else None + for item in values: + captured = _find_runtime_value(item) + if captured is not None: + return captured + return None + + __all__ = [ "FuncSpec", "FuncTemplate", diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index e3f5e4b8c8..8a52755274 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -53,6 +53,7 @@ IntegerType, Operation, StringAttr, + Type, UnitAttr, ) @@ -1011,13 +1012,21 @@ def lookup_helper(self, symbol_name: str): def _normalize_ptodsl_func_argument(self, name: str, annotation, value): raw_value = unwrap_surface_value(value) - if hasattr(raw_value, "type"): - return raw_value + target_type = None if annotation is not inspect.Parameter.empty: try: target_type = _resolve(annotation) except Exception: target_type = annotation + if hasattr(raw_value, "type"): + if not isinstance(target_type, Type): + return raw_value + return coerce_scalar_to_type( + raw_value, + target_type, + context=f"@pto.func parameter {name!r}", + ) + if target_type is not None: try: return coerce_scalar_to_type( raw_value, diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 224934a40b..975f0d827a 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1639,6 +1639,11 @@ def func_void_helper(): pto.pipe_barrier(pto.Pipe.ALL) +@pto.func(returns=None) +def func_i32_argument_helper(value: pto.i32): + _ = value + + @pto.func(returns=pto.i32) def func_partition_metadata_helper(part: pto.PartitionTensorView, cols: pto.i32): return part.sizes[0] + cols @@ -1729,6 +1734,25 @@ def ptodsl_func_constexpr_runtime_value_probe(rows: pto.i32): _ = func_constexpr_static_helper(rows, BLOCK=rows) +@pto.jit(target="a5") +def ptodsl_func_traced_argument_coercion_probe(value: pto.i64): + func_i32_argument_helper(value) + + +@pto.jit(target="a5") +def ptodsl_func_traced_argument_type_error_probe(value: pto.f32): + func_i32_argument_helper(value) + + +@pto.jit(target="a5") +def ptodsl_func_runtime_closure_capture_probe(value: pto.i32): + @pto.func(returns=None) + def helper(): + _ = value + + helper() + + @pto.jit(target="a5") def ptodsl_func_if_early_return_probe(rows: pto.i32): init = pto.const(0, dtype=pto.i32) @@ -6453,6 +6477,32 @@ def _enter_inline_simt_with_resource_attr(): lambda: ptodsl_func_constexpr_runtime_value_probe.compile().mlir_text(), "const_expr parameter 'BLOCK' expects a compile-time Python value", ) + ptodsl_func_traced_argument_coercion_text = ( + ptodsl_func_traced_argument_coercion_probe.compile().mlir_text() + ) + expect_parse_roundtrip_and_verify( + ptodsl_func_traced_argument_coercion_text, + "@pto.func traced argument annotation coercion", + ) + expect( + "arith.trunci" in ptodsl_func_traced_argument_coercion_text + and re.search( + r"func\.func @func_i32_argument_helper__ptodsl_[0-9a-f]+\(%arg0: i32\)", + ptodsl_func_traced_argument_coercion_text, + ) + is not None, + "@pto.func should adapt traced SSA arguments to their declared parameter type", + ) + expect_raises( + TypeError, + lambda: ptodsl_func_traced_argument_type_error_probe.compile().mlir_text(), + "@pto.func parameter 'value' cannot coerce", + ) + expect_raises( + TypeError, + lambda: ptodsl_func_runtime_closure_capture_probe.compile().mlir_text(), + "captures runtime value 'value'", + ) expect_raises( PTODSLAstRewriteError, lambda: ptodsl_func_if_early_return_probe.compile().mlir_text(), From 7e79d8387539d8bd0a5484736d501baad4a9554f Mon Sep 17 00:00:00 2001 From: andodo Date: Mon, 10 Aug 2026 17:40:00 +0800 Subject: [PATCH 119/122] Stabilize pto.dtype cache signatures --- ptodsl/ptodsl/_types.py | 4 +++ ptodsl/tests/test_jit_compile.py | 58 ++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/ptodsl/ptodsl/_types.py b/ptodsl/ptodsl/_types.py index d44d0f93bb..ac5677234a 100644 --- a/ptodsl/ptodsl/_types.py +++ b/ptodsl/ptodsl/_types.py @@ -79,6 +79,10 @@ def __call__(self, value): return _materialize_integer_literal(target_type, value) raise TypeError(f"unsupported eager constructor target type {target_type}") + def __ptodsl_cache_signature__(self): + with make_context(): + return ("dtype", str(self.resolve())) + def __repr__(self): return f"" diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 975f0d827a..d602f82f3e 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -11,6 +11,7 @@ from pathlib import Path import os import re +import subprocess import sys from tempfile import TemporaryDirectory from importlib.util import module_from_spec, spec_from_file_location @@ -23,6 +24,7 @@ import ptodsl._vmi_namespace as vmi_namespace from ptodsl._ast_rewrite import PTODSLAstRewriteError from ptodsl._context import make_context +from ptodsl._cache_signature import cache_signature_atom from ptodsl._kernel_signature import DeviceParameterSpec, HelperMarkerParameterSpec, RuntimeScalarParameterSpec from ptodsl._tracing.runtime import SignatureTracingRuntime from ptodsl._runtime import native_build as native_build_runtime @@ -75,6 +77,43 @@ def mlir_op_sequence(text: str) -> list[str]: return ops +def run_python_snippet(source: str) -> str: + env = dict(os.environ) + if "PYTHONPYCACHEPREFIX" not in env: + env["PYTHONPYCACHEPREFIX"] = "/tmp/ptoas-pycache" + result = subprocess.run( + [sys.executable, "-c", source], + check=True, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + return result.stdout.strip() + + +def ptodsl_func_stable_symbol_from_subprocess() -> str: + source = r''' +import re +from ptodsl import pto + +@pto.func(returns=pto.i32) +def stable_dtype_return_helper(x: pto.i32): + return x + +@pto.jit(target="a5") +def stable_dtype_symbol_probe(x: pto.i32): + _ = stable_dtype_return_helper(x) + +text = stable_dtype_symbol_probe.compile().mlir_text() +match = re.search(r"func\.func @(stable_dtype_return_helper__ptodsl_[0-9a-f]+)\(", text) +if match is None: + raise RuntimeError(text) +print(match.group(1)) +''' + return run_python_snippet(source) + + expect_raises( TypeError, lambda: pto.for_(0, 1, step=1, iter_args=(0,)), @@ -6454,6 +6493,25 @@ def _enter_inline_simt_with_resource_attr(): is not None, "@pto.func should resolve PEP 563 -> None annotations as void helpers", ) + i32_cache_signature = cache_signature_atom(pto.i32) + i64_cache_signature = cache_signature_atom(pto.i64) + ptr_cache_signature = cache_signature_atom(pto.ptr(pto.f32, "gm")) + expect( + i32_cache_signature != i64_cache_signature, + "dtype cache signatures should distinguish different MLIR scalar types", + ) + expect( + "0x" not in repr(i32_cache_signature) + and "0x" not in repr(ptr_cache_signature) + and "function" not in repr(i32_cache_signature), + "dtype cache signatures should not include Python factory reprs or memory addresses", + ) + first_stable_symbol = ptodsl_func_stable_symbol_from_subprocess() + second_stable_symbol = ptodsl_func_stable_symbol_from_subprocess() + expect( + first_stable_symbol == second_stable_symbol, + "@pto.func helper symbols should stay stable across Python processes", + ) ptodsl_func_constexpr_text = ptodsl_func_constexpr_probe.compile().mlir_text() expect_parse_roundtrip_and_verify( ptodsl_func_constexpr_text, From ea3e88e92b5efec5ef4b641c367be1669331cf21 Mon Sep 17 00:00:00 2001 From: andodo Date: Wed, 12 Aug 2026 15:12:29 +0800 Subject: [PATCH 120/122] Remove temporary pto.func probe artifacts --- tmp/ptodsl_func_chain_probe.mlir | 85 ------------------------------ tmp/ptodsl_func_chain_probe.py | 89 -------------------------------- 2 files changed, 174 deletions(-) delete mode 100644 tmp/ptodsl_func_chain_probe.mlir delete mode 100644 tmp/ptodsl_func_chain_probe.py diff --git a/tmp/ptodsl_func_chain_probe.mlir b/tmp/ptodsl_func_chain_probe.mlir deleted file mode 100644 index 4bcc3d85ed..0000000000 --- a/tmp/ptodsl_func_chain_probe.mlir +++ /dev/null @@ -1,85 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. -// Generated from tmp/ptodsl_func_chain_probe.py while developing issue #946. -// The __ptodsl_ suffixes are specialization hashes and may differ across runs. -// -// Observed behavior: -// 1. Plain undecorated helper: the internal `if True + range(2)` executes during tracing. -// The caller contains two expanded `pto.barrier ` ops and no helper func. -// 2. @pto.func dynamic loop: `for _ in range(limit)` is AST-rewritten and the helper body contains `scf.for`. -// 3. @pto.func dynamic if: `if lhs > rhs` is AST-rewritten and the helper body contains `scf.if`. -// 4. @pto.func(ast_rewrite=False): static `if True + range(2)` does not generate scf. -// It trace-time expands inside the helper body into two `arith.addi` ops. -// 5. Chained calls: `multi_return_helper -> chain_mid -> dyn_loop_helper/dyn_if_helper/no_rewrite_static_helper`. -// 6. Multiple returns: `multi_return_helper` returns `(i32, i32)`. -// 7. Reuse: `dyn_if_helper` is defined once and called twice. -// -// Counts: -// scf.for: 1 -// scf.if: 1 -// func.func @dyn_loop_helper__ptodsl_: 1 -// func.func @dyn_if_helper__ptodsl_: 1 -// func.func @no_rewrite_static_helper__ptodsl_: 1 -// func.func @chain_mid__ptodsl_: 1 -// func.func @multi_return_helper__ptodsl_: 1 -// call @dyn_if_helper__ptodsl_: 2 -// pto.barrier : 2 - -module attributes {pto.target_arch = "a5"} { - module attributes {pto.backend = "vpto", pto.kernel_kind = #pto.kernel_kind, pto.target_arch = "a5"} { - func.func @func_chain_probe(%arg0: i32) attributes {pto.entry} { - %c0_i32 = arith.constant 0 : i32 - %0:2 = call @multi_return_helper__ptodsl_c1d36bffde(%arg0, %c0_i32) : (i32, i32) -> (i32, i32) - %1 = call @dyn_if_helper__ptodsl_e04f1ff14e(%0#0, %0#1) : (i32, i32) -> i32 - pto.barrier - pto.barrier - return - } - func.func @multi_return_helper__ptodsl_c1d36bffde(%arg0: i32, %arg1: i32) -> (i32, i32) attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "multi_return_helper"} { - %0 = call @chain_mid__ptodsl_f1c09ab46a(%arg0, %arg1) : (i32, i32) -> i32 - %c1_i32 = arith.constant 1 : i32 - %1 = arith.addi %0, %c1_i32 : i32 - return %0, %1 : i32, i32 - } - func.func @chain_mid__ptodsl_f1c09ab46a(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "chain_mid"} { - %0 = call @dyn_loop_helper__ptodsl_41121a518e(%arg0, %arg1) : (i32, i32) -> i32 - %1 = call @dyn_if_helper__ptodsl_e04f1ff14e(%0, %arg1) : (i32, i32) -> i32 - %2 = call @no_rewrite_static_helper__ptodsl_3485521ccf(%1) : (i32) -> i32 - return %2 : i32 - } - func.func @dyn_loop_helper__ptodsl_41121a518e(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_loop_helper"} { - %c1_i32 = arith.constant 1 : i32 - %c0 = arith.constant 0 : index - %0 = arith.index_cast %arg0 : i32 to index - %c1 = arith.constant 1 : index - %1 = scf.for %arg2 = %c0 to %0 step %c1 iter_args(%arg3 = %arg1) -> (i32) { - %2 = arith.addi %arg3, %c1_i32 : i32 - scf.yield %2 : i32 - } - return %1 : i32 - } - func.func @dyn_if_helper__ptodsl_e04f1ff14e(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_if_helper"} { - %0 = arith.cmpi sgt, %arg0, %arg1 : i32 - %1 = scf.if %0 -> (i32) { - %2 = arith.subi %arg0, %arg1 : i32 - scf.yield %2 : i32 - } else { - %2 = arith.subi %arg1, %arg0 : i32 - scf.yield %2 : i32 - } - return %1 : i32 - } - func.func @no_rewrite_static_helper__ptodsl_3485521ccf(%arg0: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "no_rewrite_static_helper"} { - %c1_i32 = arith.constant 1 : i32 - %0 = arith.addi %arg0, %c1_i32 : i32 - %c1_i32_0 = arith.constant 1 : i32 - %1 = arith.addi %0, %c1_i32_0 : i32 - return %1 : i32 - } - } -} diff --git a/tmp/ptodsl_func_chain_probe.py b/tmp/ptodsl_func_chain_probe.py deleted file mode 100644 index d47af6e29e..0000000000 --- a/tmp/ptodsl_func_chain_probe.py +++ /dev/null @@ -1,89 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) 2026 Huawei Technologies Co., Ltd. -# This program is free software, you can redistribute it and/or modify it under the terms and conditions of -# CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -# See LICENSE in the root of the software repository for the full text of the License. -"""Probe for mixed ``@pto.func`` AST rewrite and trace-time expansion behavior.""" - -from ptodsl import pto - - -def plain_trace_helper(): - if True: - for _ in range(2): - pto.pipe_barrier(pto.Pipe.ALL) - - -@pto.func(returns=pto.i32) -def dyn_loop_helper(limit: pto.i32, value: pto.i32): - one = pto.const(1, dtype=pto.i32) - total = value - for _ in range(limit): - total = total + one - return total - - -@pto.func(returns=pto.i32) -def dyn_if_helper(lhs: pto.i32, rhs: pto.i32): - if lhs > rhs: - chosen = lhs - rhs - else: - chosen = rhs - lhs - return chosen - - -@pto.func(ast_rewrite=False, returns=pto.i32) -def no_rewrite_static_helper(value: pto.i32): - total = value - if True: - for _ in range(2): - total = total + pto.const(1, dtype=pto.i32) - return total - - -@pto.func(returns=pto.i32) -def chain_mid(limit: pto.i32, seed: pto.i32): - looped = dyn_loop_helper(limit, seed) - branched = dyn_if_helper(looped, seed) - static_expanded = no_rewrite_static_helper(branched) - return static_expanded - - -@pto.func(returns=(pto.i32, pto.i32)) -def multi_return_helper(limit: pto.i32, seed: pto.i32): - value = chain_mid(limit, seed) - return value, value + pto.const(1, dtype=pto.i32) - - -@pto.jit(target="a5") -def func_chain_probe(limit: pto.i32): - zero = pto.const(0, dtype=pto.i32) - first, second = multi_return_helper(limit, zero) - merged = dyn_if_helper(first, second) - _ = merged - plain_trace_helper() - - -def main(): - text = func_chain_probe.compile().mlir_text() - print(text) - print("\n=== COUNTS ===") - for needle in [ - "scf.for", - "scf.if", - "func.func @dyn_loop_helper__ptodsl_", - "func.func @dyn_if_helper__ptodsl_", - "func.func @no_rewrite_static_helper__ptodsl_", - "func.func @chain_mid__ptodsl_", - "func.func @multi_return_helper__ptodsl_", - "call @dyn_if_helper__ptodsl_", - "pto.barrier ", - ]: - print(f"{needle}: {text.count(needle)}") - - -if __name__ == "__main__": - main() From 83c90bbd3bba5ad076a5c1e88b21793c7511ab42 Mon Sep 17 00:00:00 2001 From: andodo Date: Wed, 12 Aug 2026 15:29:03 +0800 Subject: [PATCH 121/122] Update control-flow helper decorator docs --- ptodsl/docs/user_guide/05-control-flow.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ptodsl/docs/user_guide/05-control-flow.md b/ptodsl/docs/user_guide/05-control-flow.md index 3706cfb6fd..386444c4aa 100644 --- a/ptodsl/docs/user_guide/05-control-flow.md +++ b/ptodsl/docs/user_guide/05-control-flow.md @@ -255,7 +255,7 @@ This lets you write a single kernel that specializes into different strategies b ## 5.5 Native Python control-flow rewrite -`@pto.jit`, `@pto.func`, and named `@pto.cube` / `@pto.simd` / `@pto.simt` +`@pto.jit`, `@pto.func`, and named `@pto.tileop` / `@pto.simt` callables rewrite supported native Python control flow before tracing their bodies. In the default mode, plain Python `if` and `for range(...)` in the rewritten scope become device-side control flow. Use `pto.const_expr(...)` and From 51a667736288a6ce43288143c9b07eebe9d965ef Mon Sep 17 00:00:00 2001 From: andodo Date: Wed, 12 Aug 2026 18:03:31 +0800 Subject: [PATCH 122/122] Fix PTODSL rebase regressions --- ptodsl/ptodsl/_ast_rewrite.py | 8 ++++---- ptodsl/ptodsl/_types.py | 2 ++ 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index 6fe12b6b30..e065b8357c 100644 --- a/ptodsl/ptodsl/_ast_rewrite.py +++ b/ptodsl/ptodsl/_ast_rewrite.py @@ -34,9 +34,9 @@ def rewrite_jit_function( the optional control-flow rewrite is disabled. This keeps Python's function-local assignment rules from leaking a section-local SSA value into a sibling physical section. - ``reject_bare_returns`` controls whether a bare ``return`` inside - rewritten control flow is rejected, while value returns are always - rejected because rewritten branches must communicate through locals. + ``reject_bare_returns`` controls whether ``return`` inside rewritten + control flow is rejected. ``@pto.jit`` keeps the historical behavior, while + ``@pto.func`` enables this because helper bodies must keep one helper ABI. """ try: source = inspect.getsource(fn) @@ -1023,7 +1023,7 @@ def __init__(self, *, reject_bare_returns: bool): self._reject_bare_returns = reject_bare_returns def visit_Return(self, node): - if self._reject_bare_returns or node.value is not None: + if self._reject_bare_returns: self.exit_node = node def visit_Yield(self, node): diff --git a/ptodsl/ptodsl/_types.py b/ptodsl/ptodsl/_types.py index ac5677234a..70dd2d56f5 100644 --- a/ptodsl/ptodsl/_types.py +++ b/ptodsl/ptodsl/_types.py @@ -36,6 +36,8 @@ def softmax(arg0: pto.ptr(pto.float32, "GM"), ...): VectorType, ) +from ._context import make_context + # ── Address-space name → AddressSpace enum ─────────────────────────────────── _ADDR_SPACE = { "ub": _pto.AddressSpace.VEC, # UB == unified buffer == VEC in PTO