AGENTS.md

1# AI Assistant Guidelines for OpenXLA Development
2 
3This document provides guidelines for AI code assistants when generating,
4suggesting, or modifying code within the
5`third_party/tensorflow/compiler/xla` (OpenXLA) codebase.
6 
7## General Context
8 
9*   **Impact:** OpenXLA is a core compiler for machine learning acceleration. Changes here affect the open-source community and various hardware backends.
10*   **Code Quality:** Adhere to [Google C++ Style Guide](https://google.github.io/styleguide/cppguide.html) and OpenXLA-specific conventions.
11*   **Portability:** This code is open-sourced and runs on a number of host platforms (e.g. Linux, Windows, etc.)
12 
13## Coding Guidelines for AI Assistance
14 
151.  **Error Handling (`absl::Status`, `absl::StatusOr`)**:
16    *   **Always** use `absl::Status` or `absl::StatusOr<T>` for functions that can encounter recoverable errors.
17    *   **Macros**:
18        *   Use header `tsl/platform/status_macros.h`.
19        *   Use `RETURN_IF_ERROR` for error propagation.
20        *   Use `ASSIGN_OR_RETURN` for `StatusOr` assignments.
21    *   **Safely access `StatusOr<T>` values**: Check `.ok()` before accessing.
22 
232.  **Assertions & Invariant Checks (`TF_RET_CHECK`)**:
24    *   **Avoid `DCHECK` / `LOG(DFATAL)`** for checking returnable errors.
25    *   **Prefer `TF_RET_CHECK`**:
26        *   Located in `xla/status_macros.h`.
27        *   Use strict internal invariant checks inside functions returning `absl::Status` / `StatusOr`.
28        *   Example:
29```cpp
30            #include "xla/status_macros.h"
31 
32            absl::Status Process(const Thing* t) {
33              TF_RET_CHECK(t != nullptr) << "Thing cannot be null";
34              // ...
35              return absl::OkStatus();
36            }
37```
383.  **Decision making**:
39    *   **Avoid `bool`** for returning a decision to do or not to do something.
40    *   **Prefer `Decision`**:
41        *   Located in `third_party/tensorflow/compiler/xla/service/decision.h`
42        *   Example:
43```cpp
44            #include "third_party/tensorflow/compiler/xla/service/decision.h"
45 
46            using AutotunerDecision = Decision;
47 
48            AutotunerDecision ShouldAutotuneCublasCall(HloInstruction* instr) {
49                // ...
50                return AutotunerDecision::Forbid("Cublas autotuning was explicitly disabled");
51            }
52 
53            voud AutotuneCublas(const AutotunerDecision& decision) {
54                if (decision.IsForbidden()) {
55                    return;
56                }
57                ...
58```
59 
604.  **Performance Sensitivity**:
61    *   OpenXLA is a compiler; patterns should be efficient.
62    *   Avoid unnecessary string copies or expensive allocations in hot paths (e.g., HLO passes).
63 
645.  **Testing**:
65    *   Write unit tests using `EXPECT_EQ`, `EXPECT_TRUE`, etc.
66    *   Use macros in `tsl/platform/status_matchers.h` instead of their TF_*
67        counterparts. For example:
68        *   Use `ASSERT_OK_AND_ASSIGN`, `ASSERT_OK`, and `EXPECT_OK`.
69        *   DO NOT USE `TF_ASSERT_OK_AND_ASSIGN`, `TF_ASSERT_OK`, and
70            `TF_EXPECT_OK`.
71        *   When you refactor code that uses the TF_* macros., replace them, but
72            do not touch unrelated code.
73    *   Put tests into an anonymous namespace
74    *   Use `HloPjRtInterpreterReferenceMixin<HloPjRtTestBase>` or
75        `HloHardwareIndependentTestBase` for compiler pass tests locally where
76        possible.
77    *   Ensure tests are deterministic and do not flake.
78 
796.  **BUILD targets**:
80    * When defining BUILD targets prefer these XLA specific rules:
81        *   Instead of `proto_library` use `tf_proto_library`. There is no need
82            to define language specific targets with `tf_proto_library`.
83        *   Instead of `cc_test` use `xla_cc_test`.
84 
857.  **Explicit Typing**:
86    *   **Avoid `auto`** in public headers or complex logic chains.
87 
888.  **Compiler Phases & Invariants**:
89    *   **Phase Ordering**: Understand where your pass or change sits in the pipeline (e.g., Optimizations, Layout Assignment, Fusion).
90    *   **Invariants**: Respect the invariants of the current phase.
91        *   *Example*: Do not generate `kCustomCall` instructions before the relevant expansion pass if they are not supported by the HLO verifier at that stage.
92        *   *Example*: Do not rely on layout information before Layout Assignment.
93 
949.  **Namespaces**:
95    *   Prefer xla::gpu over nested namespaces.
96 
9710. **MLIR Operation Creation**:
98    *   **Always** use the static `OpTy::create(rewriter, ...)` method when creating MLIR operations.
99    *   **Avoid** using `rewriter.create<OpTy>(...)`. This syntax is deprecated.
100