add enzyme_noatomic attribute by minansys · Pull Request #2661 · EnzymeAD/Enzyme

minansys · 2026-01-15T22:47:30Z

Preserve and honor enzyme_noatomic on CUDA functions/args

Add support for the "enzyme_noatomic" annotation in PreserveNVVM and
teach the reverse-mode atomic path to disable atomics when either the
function or a base pointer argument carries that attribute. This keeps
CUDA annotations round-tripped into LLVM attributes and applies them
consistently during gradient generation.

Usage: mark a kernel or helper with attribute((annotate("enzyme_noatomic")))
(or attach "enzyme_noatomic" to a specific pointer argument), run
preserve-nvvm before Enzyme, and Enzyme will avoid atomicrmw in the
generated reverse kernel.


// Minimal CUDA kernel to exercise enzyme_noatomic via annotations.
#include <cuda_runtime.h>
#define ENZYME_NOATOMIC __attribute__((annotate("enzyme_noatomic")))
extern "C" {
__device__ int enzyme_dup;
__device__ void __enzyme_autodiff(void *, ...);
}
extern "C" __device__ ENZYME_NOATOMIC void my_kernel_body(float *out,
                                                          const float *in) {
  uint tid = blockIdx.x * blockDim.x + threadIdx.x;
  out[0] += in[tid];
}
extern "C" __device__ void test_derivative(float *out, float *dout,
                                           const float *in, float *din) {
  __enzyme_autodiff((void *)my_kernel_body, enzyme_dup, out, dout, enzyme_dup,
                    in, din);
}

the result with and without the ENZYME_NOATOMIC

Copilot

Pull request overview

Adds a new Enzyme annotation/attribute plumbing path intended to influence reverse-mode accumulation on GPU targets by disabling atomic updates when a marker attribute is present, and preserves that marker through the NVVM preservation pass.

Changes:

Teach PreserveNVVM to convert the llvm.global.annotations string "enzyme_elementwise_read" into an LLVM function attribute.
Add a Clang plugin parsed attribute enzyme_elementwise_read that lowers to AnnotateAttr("enzyme_elementwise_read").
Disable reverse-mode atomic accumulation in DiffeGradientUtils::addToInvertedPtrDiffe when the function (and nominally a base-pointer argument) carries "enzyme_elementwise_read".

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
`enzyme/Enzyme/PreserveNVVM.cpp`	Preserves a new annotation string by translating it into an LLVM function attribute.
`enzyme/Enzyme/DiffeGradientUtils.cpp`	Adds an attribute check that can disable atomic derivative accumulation based on context.
`enzyme/Enzyme/Clang/EnzymeClang.cpp`	Adds Clang-side parsing for `enzyme_elementwise_read` and lowers it to `AnnotateAttr`.

enzyme/Enzyme/PreserveNVVM.cpp

enzyme/Enzyme/DiffeGradientUtils.cpp

enzyme/Enzyme/Clang/EnzymeClang.cpp

Copilot · 2026-04-09T20:12:59Z

enzyme/Enzyme/DiffeGradientUtils.cpp

+  if (Atomic && elementwiseReadForContext(orig, origptr))
+    Atomic = false;


This change alters reverse-mode accumulation behavior by disabling atomics when the new attribute is present, but there’s no lit/IR test exercising the new attribute path (e.g., ensuring PreserveNVVM round-trips the annotation into an LLVM fn attribute and that the generated reverse kernel avoids atomicrmw for that accumulation). Please add a regression test covering both the attribute preservation and the expected removal of atomicrmw in the differentiated output.

@copilot apply changes based on this feedback

minansys · 2026-04-11T13:44:58Z

@copilot fix this clang format issue "--- enzyme/Enzyme/PreserveNVVM.cpp (original)
+++ enzyme/Enzyme/PreserveNVVM.cpp (reformatted)
@@ -388,10 +388,9 @@
}

         if (AS == "enzyme_elementwise_read" && Func) {

```
         Func->addAttribute(
```

             AttributeList::FunctionIndex,

             Attribute::get(Func->getContext(),

                            "enzyme_elementwise_read"));

         Func->addAttribute(AttributeList::FunctionIndex,

                            Attribute::get(Func->getContext(),

                                           "enzyme_elementwise_read"));
         changed = true;
         replacements.push_back(Constant::getNullValue(CAOp->getType()));
         continue;"

minansys · 2026-04-13T13:28:08Z

@copilot fix the unit test failure due to the new added ll test "RUN: at line 1: if [ 15 -lt 16 ]; then /opt/homebrew/Cellar/llvm@15/15.0.7/bin/opt < /Users/runner/work/Enzyme/Enzyme/enzyme/test/Enzyme/ReverseMode/elementwise-read.ll --enable-new-pm=0 -load=/Users/runner/work/Enzyme/Enzyme/build/Enzyme/LLVMEnzyme-15.dylib --enzyme-attributor=0 -enzyme-preopt=false -preserve-nvvm -enzyme -enzyme-detect-readthrow=0 -S | FileCheck /Users/runner/work/Enzyme/Enzyme/enzyme/test/Enzyme/ReverseMode/elementwise-read.ll; fi

'[' 15 -lt 16 ']'
/opt/homebrew/Cellar/llvm@15/15.0.7/bin/opt --enable-new-pm=0 -load=/Users/runner/work/Enzyme/Enzyme/build/Enzyme/LLVMEnzyme-15.dylib --enzyme-attributor=0 -enzyme-preopt=false -preserve-nvvm -enzyme -enzyme-detect-readthrow=0 -S
FileCheck /Users/runner/work/Enzyme/Enzyme/enzyme/test/Enzyme/ReverseMode/elementwise-read.ll
/Users/runner/work/Enzyme/Enzyme/enzyme/test/Enzyme/ReverseMode/elementwise-read.ll:29:3: error: found 'CHECK-LABEL:' with variable definition or use
; CHECK-LABEL: define float @VMul(float addrspace(1)* %inp) #[[VMULATTR:[0-9]+]]"

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

Copilot · 2026-04-13T13:31:49Z

enzyme/Enzyme/PreserveNVVM.cpp

+            if (AS == "enzyme_elementwise_read" && Func) {
+              Func->addAttribute(AttributeList::FunctionIndex,
+                                 Attribute::get(Func->getContext(),
+                                                "enzyme_elementwise_read"));
+              changed = true;
+              replacements.push_back(Constant::getNullValue(CAOp->getType()));


The PR title/description refer to adding/preserving an "enzyme_noatomic" annotation, but this change implements preservation of "enzyme_elementwise_read" (and there are no occurrences of "enzyme_noatomic" in the code/tests). Please clarify whether the intended attribute name is "enzyme_noatomic" or "enzyme_elementwise_read", and align the implementation + tests + documentation accordingly.

Copilot · 2026-04-13T13:31:49Z

enzyme/Enzyme/Clang/EnzymeClang.cpp

+  bool diagAppertainsToDecl(Sema &S, const ParsedAttr &Attr,
+                            const Decl *D) const override {
+    if (isa<FunctionDecl>(D))
+      return true;
+    S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
+        << Attr << "functions";
+    return false;


This adds support for applying the attribute via the Clang plugin, but the handler only appertains to FunctionDecl. The PR description mentions attaching the annotation to specific pointer arguments; if argument-level marking is required, this should also be permitted on ParmVarDecl (and lowered to an LLVM param attribute or another representation that Enzyme can detect).

Copilot · 2026-04-13T13:31:50Z

enzyme/Enzyme/DiffeGradientUtils.cpp

  // all additional parallelism in this function is outlined.
  if (backwardsOnlyShadows.find(TmpOrig) != backwardsOnlyShadows.end())
    Atomic = false;
+  if (Atomic && elementwiseReadForContext(orig, origptr))
+    Atomic = false;


addToInvertedPtrDiffe already computes TmpOrig = getBaseObject(origptr) just above, but elementwiseReadForContext calls getBaseObject(origptr) again. Since this code is on a hot path, consider reusing the already-computed base object (e.g., pass TmpOrig into the helper or have the helper accept the base Value*) to avoid duplicate value-tracking work.

minansys and others added 3 commits January 15, 2026 17:43

add enzyme_noatomic attribute

dc07e73

update

cab5a3a

Merge branch 'main' into mixu/noatomic_attribute

20d65fc

minansys requested a review from Copilot April 9, 2026 20:09

Copilot started reviewing on behalf of minansys April 9, 2026 20:10 View session

Copilot AI reviewed Apr 9, 2026

View reviewed changes

add test and fix the clang format

ad226a9

minansys requested a review from Copilot April 13, 2026 13:28

Copilot started reviewing on behalf of minansys April 13, 2026 13:29 View session

Copilot AI reviewed Apr 13, 2026

View reviewed changes

minansys added 2 commits April 13, 2026 09:54

fix the unit tests

5a5930a

fix ci test

26ff8b5

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

add enzyme_noatomic attribute#2661

add enzyme_noatomic attribute#2661
minansys wants to merge 6 commits intoEnzymeAD:mainfrom
minansys:mixu/noatomic_attribute

minansys commented Jan 15, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Apr 9, 2026

Uh oh!

minansys Apr 11, 2026

Uh oh!

minansys commented Apr 11, 2026

Uh oh!

minansys commented Apr 13, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Apr 13, 2026

Uh oh!

Copilot AI Apr 13, 2026

Uh oh!

Copilot AI Apr 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		if (Atomic && elementwiseReadForContext(orig, origptr))
		Atomic = false;

Conversation

minansys commented Jan 15, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

minansys Apr 11, 2026

Choose a reason for hiding this comment

Uh oh!

minansys commented Apr 11, 2026

Uh oh!

minansys commented Apr 13, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Apr 13, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 13, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 13, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants