Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Batched autodiff #2181

Merged
merged 22 commits into from
Dec 27, 2024
Merged

Batched autodiff #2181

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
8d3ef08
add type conversions for width != 1.
jumerckx Nov 28, 2024
5e647e5
add code to tblgen generator, this eventually needs to be a single fu…
jumerckx Nov 28, 2024
f69f974
a test and formatting
jumerckx Nov 28, 2024
623fccf
use tensor splatop
jumerckx Dec 2, 2024
f860e19
remove stale enzyme-tblgen changes
jumerckx Dec 2, 2024
3a3bdf8
do the simple batching in enzyme-tblgen
jumerckx Dec 4, 2024
f1b2a6d
include tensor in all AutoDiffOpInterfaceImpls
jumerckx Dec 4, 2024
396a29b
add enzyme broadcastop
jumerckx Dec 6, 2024
64489a5
getShadowType for TensorTypeInterface
jumerckx Dec 6, 2024
0ccb37f
create broadcastop in enzyme-tblgen
jumerckx Dec 6, 2024
123a985
Revert "include tensor in all AutoDiffOpInterfaceImpls"
jumerckx Dec 6, 2024
3d2911d
test
jumerckx Dec 6, 2024
8a3f3c7
DenseI64ArrayAttr for shape instead of scalar width
Dec 22, 2024
aa4e598
`llvm::SmallVector` --> `ArrayRef`
jumerckx Dec 22, 2024
fa190a4
formatting
jumerckx Dec 23, 2024
1802949
use getShadowType in BroadcastOp builder
jumerckx Dec 25, 2024
f9746ac
unstructured control flow test
jumerckx Dec 25, 2024
2fb1412
scf.for
jumerckx Dec 25, 2024
2c4f1fb
formatting
jumerckx Dec 25, 2024
236b65d
support `scf.if` test
jumerckx Dec 25, 2024
9ceb922
formatting
jumerckx Dec 25, 2024
50522a0
forgotten includes
jumerckx Dec 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
scf.for
  • Loading branch information
@jumerckx
jumerckx committed Dec 25, 2024
commit 2fb1412d81a4798d95eaba6a6640c1f7e60d6039
2 changes: 1 addition & 1 deletion enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ mlir::Value mlir::enzyme::MGradientUtils::invertPointerM(mlir::Value v,
return invertedPointers.lookupOrNull(v);

if (isConstantValue(v)) {
if (auto iface = v.getType().dyn_cast<AutoDiffTypeInterface>()) {
if (auto iface = getShadowType(v.getType()).dyn_cast<AutoDiffTypeInterface>()) {
OpBuilder::InsertionGuard guard(Builder2);
if (auto op = v.getDefiningOp())
Builder2.setInsertionPoint(getNewFromOriginal(op));
Expand Down
33 changes: 33 additions & 0 deletions enzyme/test/MLIR/ForwardMode/batched_for.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
// RUN: %eopt --enzyme %s | FileCheck %s

module {
func.func @square(%x : f64) -> f64 {
%cst = arith.constant 10.000000e+00 : f64
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c10 = arith.constant 10 : index
%r = scf.for %arg1 = %c0 to %c10 step %c1 iter_args(%arg2 = %cst) -> (f64) {
%n = arith.addf %arg2, %x : f64
scf.yield %n : f64
}
return %r : f64
}
func.func @dsq(%x : f64, %dx : tensor<2xf64>) -> tensor<2xf64> {
%r = enzyme.fwddiff @square(%x, %dx) { activity=[#enzyme<activity enzyme_dup>], ret_activity=[#enzyme<activity enzyme_dupnoneed>], width=2 } : (f64, tensor<2xf64>) -> (tensor<2xf64>)
return %r : tensor<2xf64>
}
}

// CHECK: func.func private @fwddiffesquare(%[[arg0:.+]]: f64, %[[arg1:.+]]: tensor<2xf64>) -> tensor<2xf64> {
// CHECK-DAG: %[[cst:.+]] = arith.constant dense<0.000000e+00> : tensor<2xf64>
// CHECK-DAG: %[[cst_0:.+]] = arith.constant 1.000000e+01 : f64
// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[c10:.+]] = arith.constant 10 : index
// CHECK-NEXT: %[[i0:.+]]:2 = scf.for %[[arg2:.+]] = %[[c0]] to %[[c10]] step %[[c1]] iter_args(%[[arg3:.+]] = %[[cst_0]], %[[arg4:.+]] = %[[cst]]) -> (f64, tensor<2xf64>) {
// CHECK-NEXT: %[[i1:.+]] = arith.addf %[[arg4]], %[[arg1]] : tensor<2xf64>
// CHECK-NEXT: %[[i2:.+]] = arith.addf %[[arg3]], %[[arg0]] : f64
// CHECK-NEXT: scf.yield %[[i2]], %[[i1]] : f64, tensor<2xf64>
// CHECK-NEXT: }
// CHECK-NEXT: return %[[i0]]#1 : tensor<2xf64>
// CHECK-NEXT: }
Loading