Skip to content

Commit e3a201f

Browse files
jumerckxjumerckx
committed
Squashed commit of the following:
commit 65e9976 Author: William Moses <[email protected]> Date: Sat Dec 14 14:05:03 2024 -0600 Interp2 (EnzymeAD#365) * WIP: kernels * more files * fix * wip * wqtmp * wip * inc * continuing * wip * more work * inf rec * fix * overload working * continuing * continuing * push * fix `call_with_reactant_generator` for Julia 1.11 (EnzymeAD#359) * conversion * continuing * Cleanup * Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Delete test/cuda.jl * fixup * Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * fix apply * indep of change * minor fix in name * Update utils.jl * Interp take 2 * continuing adentures * delcode * fix * tmp * make * fix * cleanup * continuing * more working * further simplify * fx * more improvements * minus show * less prints * even fewer * confusion * tmp * force clean * force oc * clean * Rewrite * fixup * fix * fix * fix * fixup * fix * wip * safe prints * fix * fix * stackoverflow * cleanup * dyindex * rt * continue * clean * fix * fix * fix * fix * fixup * fix * fix * capture oc * compile perf * v1.11 fix * other way 'round * formatting --------- Co-authored-by: William Moses <[email protected]> Co-authored-by: jumerckx <[email protected]> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: jumerckx <[email protected]> commit 73899f5 Author: Avik Pal <[email protected]> Date: Sat Dec 14 14:58:47 2024 +0530 fix: include files if they end with .jl (EnzymeAD#377) commit 9f96c09 Author: Sergio Sánchez Ramírez <[email protected]> Date: Fri Dec 13 23:12:43 2024 +0100 Run CI on aarch64 (EnzymeAD#350) * Run CI on aarch64 * use julia pipeline for aarch64-linux * fix var * exclude aarch64-linux jobs from github ci commit b56e661 Author: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Date: Fri Dec 13 08:58:51 2024 +0530 chore: format code (EnzymeAD#371) Co-authored-by: mofeing <[email protected]> commit 311498b Author: Anton Oresten <[email protected]> Date: Thu Dec 12 05:41:39 2024 +0100 feat: define outer `repeat` method for `TracedRArray` (EnzymeAD#361) * Add repeat method * Add repeat tests * Update test/basic.jl * Update src/TracedRArray.jl commit 8b90501 Author: Avik Pal <[email protected]> Date: Thu Dec 12 10:11:00 2024 +0530 fix: ensure printing of wrapped ConcreteRArrays goes through our show (EnzymeAD#367) * fix: ensure printing of wrapped ConcreteRArrays goes through our show * fix: allow wrapped arrays in mapreduce commit ea97be3 Author: Sergio Sánchez Ramírez <[email protected]> Date: Wed Dec 11 22:02:14 2024 +0100 Ignore versioned manifests
1 parent 38b1797 commit e3a201f

26 files changed

+1705
-999
lines changed

‎.buildkite/pipeline.yml‎

Lines changed: 71 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,37 +1,76 @@
11
steps:
2-
- label: "CUDA Julia v{{matrix.version}} -- {{matrix.group}}"
3-
matrix:
4-
setup:
5-
version:
6-
- "1.10"
7-
group:
8-
- core
9-
- neural_networks
10-
- integration
11-
plugins:
12-
- JuliaCI/julia#v1:
13-
version: "{{matrix.version}}"
14-
- JuliaCI/julia-coverage#v1:
15-
codecov: true
16-
dirs:
17-
- src
18-
- ext
19-
- lib/ReactantCore/src
20-
commands: |
21-
julia --project=. -e 'println("--- :julia: Instantiating project")
22-
using Pkg
23-
Pkg.develop([PackageSpec(path="lib/ReactantCore")])'
2+
- group: ":test_tube: Tests"
3+
steps:
4+
- label: "CUDA Julia v{{matrix.version}} -- {{matrix.group}}"
5+
matrix:
6+
setup:
7+
version:
8+
- "1.10"
9+
group:
10+
- core
11+
- neural_networks
12+
- integration
13+
plugins:
14+
- JuliaCI/julia#v1:
15+
version: "{{matrix.version}}"
16+
- JuliaCI/julia-coverage#v1:
17+
codecov: true
18+
dirs:
19+
- src
20+
- ext
21+
- lib/ReactantCore/src
22+
commands: |
23+
julia --project=. -e 'println("--- :julia: Instantiating project")
24+
using Pkg
25+
Pkg.develop([PackageSpec(path="lib/ReactantCore")])'
26+
27+
julia --project=. -e 'println("--- :julia: Run Tests")
28+
using Pkg
29+
Pkg.test(; coverage="user")'
30+
agents:
31+
queue: "juliagpu"
32+
cuda: "*"
33+
env:
34+
REACTANT_TEST_GROUP: "{{matrix.group}}"
35+
if: build.message !~ /\[skip tests\]/
36+
timeout_in_minutes: 60
37+
38+
- label: ":julia: :linux: aarch64 - Julia v{{matrix.version}} -- {{matrix.group}}"
39+
matrix:
40+
setup:
41+
version:
42+
- "1.10"
43+
- "1.11"
44+
group:
45+
- core
46+
- neural_networks
47+
- integration
48+
plugins:
49+
- JuliaCI/julia#v1:
50+
version: "{{matrix.version}}"
51+
- JuliaCI/julia-coverage#v1:
52+
codecov: true
53+
dirs:
54+
- src
55+
- ext
56+
- lib/ReactantCore/src
57+
commands: |
58+
julia --project=. -e 'println("--- :julia: Instantiating project")
59+
using Pkg
60+
Pkg.develop([PackageSpec(path="lib/ReactantCore")])'
2461
25-
julia --project=. -e 'println("--- :julia: Run Tests")
26-
using Pkg
27-
Pkg.test(; coverage="user")'
28-
agents:
29-
queue: "juliagpu"
30-
cuda: "*"
31-
env:
32-
REACTANT_TEST_GROUP: "{{matrix.group}}"
33-
if: build.message !~ /\[skip tests\]/
34-
timeout_in_minutes: 60
62+
julia --project=. -e 'println("--- :julia: Run Tests")
63+
using Pkg
64+
Pkg.test(; coverage="user")'
65+
agents:
66+
queue: "juliaecosystem"
67+
os: "linux"
68+
sandbox_capable: "true"
69+
arch: "aarch64"
70+
env:
71+
REACTANT_TEST_GROUP: "{{matrix.group}}"
72+
if: build.message !~ /\[skip tests\]/
73+
timeout_in_minutes: 60
3574

3675
- group: ":racehorse: Benchmarks"
3776
steps:

‎.github/workflows/CI.yml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ jobs:
3333
- integration
3434
arch:
3535
- x64
36+
- aarch64
3637
assertions:
3738
- false
3839
libReactant: [packaged]
@@ -64,6 +65,10 @@ jobs:
6465
libReactant: packaged
6566
version: '1.10'
6667
test_group: integration
68+
exclude:
69+
# these are run on Buildkite
70+
- os: ubuntu-20.04
71+
arch: aarch64
6772
steps:
6873
- uses: actions/checkout@v4
6974
- uses: julia-actions/setup-julia@v2

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,7 @@ docs/site/
251251
# committed for packages, but should be committed for applications that require a static
252252
# environment.
253253
Manifest.toml
254+
Manifest-v*.toml
254255

255256
.vscode/*
256257
.vscode/settings.json

‎Project.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,14 +42,14 @@ Adapt = "4"
4242
ArrayInterface = "7.10"
4343
CEnum = "0.4, 0.5"
4444
Downloads = "1.6"
45-
Enzyme = "0.13.21"
45+
Enzyme = "0.13.22"
4646
EnzymeCore = "0.8.8"
4747
GPUArraysCore = "0.1.6, 0.2"
4848
LinearAlgebra = "1.10"
4949
NNlib = "0.9.26"
5050
OrderedCollections = "1"
5151
Preferences = "1.4"
52-
ReactantCore = "0.1.2"
52+
ReactantCore = "0.1.3"
5353
Reactant_jll = "0.0.26"
5454
ScopedValues = "1.2.1"
5555
Scratch = "1.2"

‎deps/ReactantExtra/API.cpp‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,16 @@ extern "C" MlirModule ConvertLLVMToMLIR(LLVMModuleRef lmod, MlirContext cctx) {
376376
return wrap(res);
377377
}
378378

379+
#include "llvm/IRReader/IRReader.h"
380+
extern "C" MlirModule ConvertLLVMStrToMLIR(const char* lmod, MlirContext cctx) {
381+
LLVMContext Context;
382+
SMDiagnostic Err;
383+
auto llvmModule = llvm::parseIR(llvm::MemoryBufferRef(lmod, "conversion"), Err, Context);
384+
mlir::MLIRContext &context = *unwrap(cctx);
385+
auto res = mlir::translateLLVMIRToModule(std::move(llvmModule), &context, /*emitExpensiveWarnings*/false, /*dropDICompositeElements*/false).release();
386+
return wrap(res);
387+
}
388+
379389

380390
/* Note that this */
381391
extern "C" xla::PjRtLoadedExecutable* ClientCompile(PjRtClient * client, MlirModule cmod) {

‎deps/ReactantExtra/BUILD‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -450,6 +450,8 @@ cc_library(
450450
"@llvm-project//mlir:SCFDialect",
451451
"@llvm-project//mlir:TransformDialect",
452452
"@llvm-project//mlir:Transforms",
453+
454+
"@llvm-project//llvm:IRReader",
453455
"@llvm-project//llvm:Support",
454456
"@llvm-project//llvm:AArch64AsmParser",
455457
"@llvm-project//llvm:AArch64CodeGen",

‎ext/ReactantNNlibExt.jl‎

Lines changed: 11 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,10 @@ module ReactantNNlibExt
22

33
using NNlib
44
using GPUArraysCore: @allowscalar
5-
using Reactant:
6-
Reactant,
7-
Ops,
8-
TracedRArray,
9-
AnyTracedRArray,
10-
materialize_traced_array,
11-
MLIR,
12-
TracedRNumber,
13-
get_mlir_data,
14-
set_mlir_data!
5+
using Reactant: Reactant, Ops, TracedRArray, AnyTracedRArray, MLIR, TracedRNumber
6+
7+
using Reactant.TracedUtils: materialize_traced_array, get_mlir_data, set_mlir_data!
8+
159
using ReactantCore: @trace
1610
using LinearAlgebra: LinearAlgebra, triu
1711

@@ -238,9 +232,9 @@ function NNlib.batched_mul!(
238232
if size(x, 3) != size(y, 3)
239233
B = max(size(x, 3), size(y, 3))
240234
if size(x, 3) == 1
241-
x = Reactant.broadcast_to_size(x, (size(x, 1), size(x, 2), B))
235+
x = Reactant.TracedUtils.broadcast_to_size(x, (size(x, 1), size(x, 2), B))
242236
elseif size(y, 3) == 1
243-
y = Reactant.broadcast_to_size(y, (size(y, 1), size(y, 2), B))
237+
y = Reactant.TracedUtils.broadcast_to_size(y, (size(y, 1), size(y, 2), B))
244238
end
245239
end
246240

@@ -250,9 +244,9 @@ function NNlib.batched_mul!(
250244
if size(x, 1) != size(y, 1)
251245
B = max(size(x, 1), size(y, 1))
252246
if size(x, 1) == 1
253-
x = Reactant.broadcast_to_size(x, (B, size(x, 2), size(x, 3)))
247+
x = Reactant.TracedUtils.broadcast_to_size(x, (B, size(x, 2), size(x, 3)))
254248
elseif size(y, 1) == 1
255-
y = Reactant.broadcast_to_size(y, (B, size(y, 2), size(y, 3)))
249+
y = Reactant.TracedUtils.broadcast_to_size(y, (B, size(y, 2), size(y, 3)))
256250
end
257251
end
258252

@@ -270,7 +264,7 @@ end
270264
function NNlib.pad_constant(
271265
x::AnyTracedRArray{T,N}, pad::NTuple{N,Tuple{Int,Int}}, value
272266
) where {T,N}
273-
value = Reactant.promote_to(TracedRNumber{T}, value)
267+
value = Reactant.TracedUtils.promote_to(TracedRNumber{T}, value)
274268
low = [i[1] for i in pad]
275269
high = [i[2] for i in pad]
276270
interior = [0 for i in pad]
@@ -329,7 +323,8 @@ function NNlib.gather!(dst::TracedRArray, src::AnyTracedRArray, idxs::AbstractAr
329323
start_sizes = ntuple(i -> size(src, i), dims)
330324
results = map(CartesianIndices(idxs)) do k
331325
res = @allowscalar src[colons..., Tuple(idxs[k])...]
332-
res isa TracedRNumber && (res = Reactant.broadcast_to_size(res, (1,)))
326+
res isa TracedRNumber &&
327+
(res = Reactant.TracedUtils.broadcast_to_size(res, (1,)))
333328
return reshape(res, start_sizes..., :)
334329
end
335330
res = reshape(cat(results...; dims=(dims + 1)), size(dst))

‎ext/ReactantStatisticsExt.jl‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
module ReactantStatisticsExt
22

3-
using Reactant: AnyTracedRArray, materialize_traced_array
3+
using Reactant: AnyTracedRArray
4+
using Reactant.TracedUtils: materialize_traced_array
45
using Statistics: Statistics
56

67
function Statistics.mean(A::AnyTracedRArray{T,N}; dims=:) where {T,N}

‎ext/ReactantYaoBlocksExt.jl‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
11
module ReactantYaoBlocksExt
22

33
using Reactant
4+
using Reactant.TracedUtils: broadcast_to_size
45
using YaoBlocks
56

67
function YaoBlocks.mat(
78
::Type{T}, R::RotationGate{D,Reactant.TracedRNumber{S},<:XGate}
89
) where {D,T,S}
9-
M = Reactant.broadcast_to_size(zero(T), (2, 2))
10+
M = broadcast_to_size(zero(T), (2, 2))
1011
c = cos(R.theta / 2)
1112
s = -im * sin(R.theta / 2)
1213
M[1, 1] = c
@@ -19,7 +20,7 @@ end
1920
function YaoBlocks.mat(
2021
::Type{T}, R::RotationGate{D,Reactant.TracedRNumber{S},<:YGate}
2122
) where {D,T,S}
22-
M = Reactant.broadcast_to_size(zero(T), (2, 2))
23+
M = broadcast_to_size(zero(T), (2, 2))
2324
c = cos(R.theta / 2)
2425
s = sin(R.theta / 2)
2526
M[1, 1] = c
@@ -32,7 +33,7 @@ end
3233
function YaoBlocks.mat(
3334
::Type{T}, R::RotationGate{D,Reactant.TracedRNumber{S},<:ZGate}
3435
) where {D,T,S}
35-
M = Reactant.broadcast_to_size(zero(T), (2, 2))
36+
M = broadcast_to_size(zero(T), (2, 2))
3637
x = exp(im * R.theta / 2)
3738
M[1, 1] = conj(x)
3839
M[2, 2] = x

‎lib/ReactantCore/Project.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "ReactantCore"
22
uuid = "a3311ec8-5e00-46d5-b541-4f83e724a433"
33
authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>"]
4-
version = "0.1.2"
4+
version = "0.1.3"
55

66
[deps]
77
ExpressionExplorer = "21656369-7473-754a-2065-74616d696c43"

0 commit comments

Comments
 (0)