Add LLVM to Affine access pass

src/enzyme_ad/jax/BUILD

@@ -390,6 +390,9 @@ cc_library(
         ":EnzymeXLAOpsIncGen",
         ":EnzymeXLAPassesIncGen",
         ":EnzymeHLOPatternsIncGen",
+        ":RaisingTransformOpsImplIncGen",
+        ":RaisingTransformOpsIncGen",
+        ":RaisingTransformPatternsIncGen",
         ":RaisingTransformOps",
         "@llvm-project//mlir:DLTIDialect",
         "@llvm-project//mlir:GPUPipelines",

src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td

@@ -98,4 +98,40 @@ def Pointer2MemrefOp : EnzymeXLA_Op<"pointer2memref", [
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// AtAddrOp
+//===----------------------------------------------------------------------===//
+
+def AtAddrOp : EnzymeXLA_Op<"ataddr", [Pure]> {
+  let summary =
+      "Construct a c-style memref at an addr";
+  let description = [{
+  }];
+  let arguments = (ins AnyType:$addr);
+  let results = (outs AnyRankedOrUnrankedMemRef:$dest);
+  let builders = [
+    OpBuilder<(ins "Value":$source), [{
+      return build($_builder, $_state,
+        TypeRange(MemRefType::get({ShapedType::kDynamic}, $_builder.getI8Type())),
+        ValueRange(source));
+    }]>,
+    OpBuilder<(ins "Value":$source, "Type":$dstType), [{
+      return build($_builder, $_state,
+        TypeRange(dstType),
+        ValueRange(source));
+    }]>
+  ];
+}
+
+def AffineScopeOp : EnzymeXLA_Op<"scope", [
+      AffineScope,
+      AutomaticAllocationScope,
+      RecursiveMemoryEffects,
+    ]>,
+    Arguments<(ins Variadic<AnyType>:$operands)>,
+    Results<(outs Variadic<AnyType>:$results)> {
+  let summary = "Inline affine scope";
+  let regions = (region SizedRegion<1>:$region);
+}
+
 #endif // ENZYMEXLA_OPS