Skip to content

Commit be3617b

Browse files
committed
autodiff: slice support in typetree
1 parent 31541fe commit be3617b

File tree

4 files changed

+36
-0
lines changed

4 files changed

+36
-0
lines changed

compiler/rustc_middle/src/ty/mod.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2327,5 +2327,12 @@ pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
23272327
}
23282328
}
23292329

2330+
if ty.is_slice() {
2331+
if let ty::Slice(element_ty) = ty.kind() {
2332+
let element_tree = typetree_from_ty(tcx, *element_ty);
2333+
return element_tree;
2334+
}
2335+
}
2336+
23302337
TypeTree::new()
23312338
}
Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
//@ needs-enzyme
2+
//@ ignore-cross-compile
3+
4+
use run_make_support::{llvm_filecheck, rfs, rustc};
5+
6+
fn main() {
7+
rustc().input("test.rs").arg("-Zautodiff=Enable").emit("llvm-ir").run();
8+
llvm_filecheck().patterns("slice.check").stdin_buf(rfs::read("test.ll")).run();
9+
}
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
; Check that slice TypeTree metadata is correctly generated
2+
; Should show Float@double for slice elements
3+
4+
CHECK: define{{.*}}"enzyme_type"="{[]:Float@double}"{{.*}}@test_slice{{.*}}"enzyme_type"="{[]:Pointer}"
Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
#![feature(autodiff)]
2+
3+
use std::autodiff::autodiff_reverse;
4+
5+
#[autodiff_reverse(d_test, Duplicated, Active)]
6+
#[no_mangle]
7+
fn test_slice(slice: &[f64]) -> f64 {
8+
slice.iter().sum()
9+
}
10+
11+
fn main() {
12+
let arr = [1.0, 2.0, 3.0, 4.0, 5.0];
13+
let slice = &arr[..];
14+
let mut d_slice = [0.0; 5];
15+
let _result = d_test(slice, &mut d_slice[..], 1.0);
16+
}

0 commit comments

Comments
 (0)