[MLIR] Add StoreLikeInterface to decouple activity analysis from stores

[vchuravy]

Activity analysis currently hard-codes which ops are stores (dyn_cast<LLVM::StoreOp> / memref::StoreOp) when reasoning about stored-value / pointer activity. That prevents out-of-tree dialects from participating and duplicates the same logic per dialect.

Change

• Add enzyme::ActiveStoreOpInterface with getStoredValue() / getStoredPointer().
• Attach it to memref.store and llvm.store.
• Use it at the potential-active-store site in ActivityAnalysis.cpp that already handled both dialects — now a single dialect-agnostic branch.

Any dialect can now opt into store-activity handling by attaching the interface (e.g. an out-of-tree fir.store / hlfir.assign), without touching the core analysis.

Behavior

Behavior-preserving: llvm.store/memref.store return the same value/pointer through the interface as the hard-coded accessors did. The remaining LLVM-specific store sites (which recurse through llvm.alloca / llvm.load origins) are intentionally left as-is.

The MLIR test suite is unchanged: 135/137 pass, and the 2 ReverseMode failures are pre-existing on this LLVM build (they fail identically without this change).

Motivation

This is a standalone extract from the Enzyme⇄Flang HLFIR work (#2969): it lets fir.store/hlfir.assign participate in activity analysis so by-reference Fortran differentiates, while keeping Analysis/ free of any FIR/HLFIR dependency.

🤖 Generated with Claude Code

[Pangoraw]

maybe could be implemented for affine.store as well?

[wsmoses]

ci err

  shell: /home/runner/externals/node20/bin/node {0}
using script executor
[0/2] Re-checking globbed directories...
[0/1] Running MLIR regression tests
-- Testing: 141 tests, 32 workers --
Testing:  0.. 10.. 20.. 
FAIL: Enzyme :: MLIR/ForwardMode/affine.mlir (44 of 141)
******************** TEST 'Enzyme :: MLIR/ForwardMode/affine.mlir' FAILED ********************
Exit Code: 1

Command Output (stderr):
--
/__w/Enzyme/Enzyme/Enzyme/enzyme-build/Enzyme/MLIR/enzymemlir-opt --enzyme /__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ForwardMode/affine.mlir | FileCheck /__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ForwardMode/affine.mlir # RUN: at line 1
+ /__w/Enzyme/Enzyme/Enzyme/enzyme-build/Enzyme/MLIR/enzymemlir-opt --enzyme /__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ForwardMode/affine.mlir
+ FileCheck /__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ForwardMode/affine.mlir
/__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ForwardMode/affine.mlir:86:16: error: CHECK-DAG: expected string not found in input
 // CHECK-DAG: %[[cst0:.+]] = arith.constant 0.000000e+00 : f64
               ^
<stdin>:80:43: note: scanning from here
 %alloc_1 = memref.alloc() : memref<1xf64>
                                          ^
<stdin>:80:43: note: pattern attempts to capture variables: "cst0"
 %alloc_1 = memref.alloc() : memref<1xf64>
                                          ^
<stdin>:86:2: note: possible intended match here
 %7 = arith.mulf %arg0, %arg0 : f64
 ^

Input file: <stdin>
Check file: /__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ForwardMode/affine.mlir

-dump-input=help explains the following input dump.

Input was:
<<<<<<
          .
          .
          .
         75:  } 
         76:  func.func private @fwddiffeif_then(%arg0: f64, %arg1: f64, %arg2: i1) -> f64 { 
         77:  %cst = arith.constant 2.000000e+00 : f64 
         78:  %cst_0 = arith.constant 1.000000e+01 : f64 
         79:  %alloc = memref.alloc() : memref<1xf64> 
         80:  %alloc_1 = memref.alloc() : memref<1xf64> 
dag:86'0                                              {   search range start (exclusive)
dag:86'1                                                  error: no match found in search range
dag:86'2                                                  pattern attempts to capture variables: "cst0"
         81:  affine.store %cst, %alloc_1[0] : memref<1xf64> 
         82:  scf.if %arg2 { 
         83:  %4 = arith.mulf %arg1, %arg0 : f64 
         84:  %5 = arith.mulf %arg1, %arg0 : f64 
         85:  %6 = arith.addf %4, %5 : f64 
         86:  %7 = arith.mulf %arg0, %arg0 : f64 
dag:86'3      ?                                    possible intended match
         87:  affine.store %6, %alloc[0] : memref<1xf64> 
         88:  affine.store %7, %alloc_1[0] : memref<1xf64> 
         89:  } 
         90:  %0 = affine.load %alloc[0] : memref<1xf64> 
         91:  %1 = affine.load %alloc_1[0] : memref<1xf64> 
         92:  %2 = arith.mulf %0, %cst : f64 
         93:  %3 = arith.mulf %cst, %1 : f64 
         94:  return %2 : f64 
         95:  } 
         96: } 
         97:  
dag:86'4      } search range end (exclusive)
>>>>>>

--

********************
Testing:  0.. 10.. 20.. 30.. 40.. 50.. 60.. 70.. 80.. 90.. 
********************
Failed Tests (1):
  Enzyme :: MLIR/ForwardMode/affine.mlir


Testing Time: 0.60s

Total Discovered Tests: 141
  Passed: 140 (99.29%)
  Failed:   1 (0.71%)
FAILED: test/MLIR/CMakeFiles/check-enzymemlir /__w/Enzyme/Enzyme/Enzyme/enzyme-build/test/MLIR/CMakeFiles/check-enzymemlir 
cd /__w/Enzyme/Enzyme/Enzyme/enzyme-build/test/MLIR && /__w/Enzyme/Enzyme/llvm-project/mlir-build/bin/llvm-lit -v -sv /__w/Enzyme/Enzyme/Enzyme/enzyme-build/test/MLIR

[wsmoses]

can you fix this xfail?

[vchuravy]

Yeah, planning on it, I just wanted to have a test first.

[wsmoses]

👍

[wsmoses]

same here