Skip to content

Commit cd8b64f

Browse files
fschlimbfschlimb
authored
Gpu compute follows data (#84)
* enabling device testing * explicitly handling return args, following new imex * generated IR letting function accept memrefs, not NDArrays * not passing local offsets as input argument * some refactoring * moving mlir business to Deferred.cpp to mlir.*D * adding jit::fini() for correct shutdown
1 parent aba6eea commit cd8b64f

28 files changed

+627
-478
lines changed

‎examples/shallow_water.py‎

Lines changed: 35 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,10 @@
2828
import time as time_mod
2929
import argparse
3030

31+
import os
32+
33+
device = os.getenv("SHARPY_USE_GPU", "")
34+
3135

3236
def run(n, backend, datatype, benchmark_mode):
3337
if backend == "sharpy":
@@ -94,16 +98,24 @@ def info(s):
9498
t_end = 1.0
9599

96100
# coordinate arrays
97-
x_t_2d = fromfunction(lambda i, j: xmin + i * dx + dx / 2, (nx, ny), dtype=dtype)
98-
y_t_2d = fromfunction(lambda i, j: ymin + j * dy + dy / 2, (nx, ny), dtype=dtype)
99-
x_u_2d = fromfunction(lambda i, j: xmin + i * dx, (nx + 1, ny), dtype=dtype)
101+
x_t_2d = fromfunction(
102+
lambda i, j: xmin + i * dx + dx / 2, (nx, ny), dtype=dtype, device=device
103+
)
104+
y_t_2d = fromfunction(
105+
lambda i, j: ymin + j * dy + dy / 2, (nx, ny), dtype=dtype, device=device
106+
)
107+
x_u_2d = fromfunction(
108+
lambda i, j: xmin + i * dx, (nx + 1, ny), dtype=dtype, device=device
109+
)
100110
y_u_2d = fromfunction(
101-
lambda i, j: ymin + j * dy + dy / 2, (nx + 1, ny), dtype=dtype
111+
lambda i, j: ymin + j * dy + dy / 2, (nx + 1, ny), dtype=dtype, device=device
102112
)
103113
x_v_2d = fromfunction(
104-
lambda i, j: xmin + i * dx + dx / 2, (nx, ny + 1), dtype=dtype
114+
lambda i, j: xmin + i * dx + dx / 2, (nx, ny + 1), dtype=dtype, device=device
115+
)
116+
y_v_2d = fromfunction(
117+
lambda i, j: ymin + j * dy, (nx, ny + 1), dtype=dtype, device=device
105118
)
106-
y_v_2d = fromfunction(lambda i, j: ymin + j * dy, (nx, ny + 1), dtype=dtype)
107119

108120
T_shape = (nx, ny)
109121
U_shape = (nx + 1, ny)
@@ -120,32 +132,32 @@ def info(s):
120132
info(f"Total DOFs: {dofs_T + dofs_U + dofs_V}")
121133

122134
# prognostic variables: elevation, (u, v) velocity
123-
e = np.full(T_shape, 0.0, dtype)
124-
u = np.full(U_shape, 0.0, dtype)
125-
v = np.full(V_shape, 0.0, dtype)
135+
e = np.full(T_shape, 0.0, dtype, device=device)
136+
u = np.full(U_shape, 0.0, dtype, device=device)
137+
v = np.full(V_shape, 0.0, dtype, device=device)
126138

127139
# potential vorticity
128-
q = np.full(F_shape, 0.0, dtype)
140+
q = np.full(F_shape, 0.0, dtype, device=device)
129141

130142
# bathymetry
131-
h = np.full(T_shape, 0.0, dtype)
143+
h = np.full(T_shape, 0.0, dtype, device=device)
132144

133-
hu = np.full(U_shape, 0.0, dtype)
134-
hv = np.full(V_shape, 0.0, dtype)
145+
hu = np.full(U_shape, 0.0, dtype, device=device)
146+
hv = np.full(V_shape, 0.0, dtype, device=device)
135147

136-
dudy = np.full(F_shape, 0.0, dtype)
137-
dvdx = np.full(F_shape, 0.0, dtype)
148+
dudy = np.full(F_shape, 0.0, dtype, device=device)
149+
dvdx = np.full(F_shape, 0.0, dtype, device=device)
138150

139151
# vector invariant form
140-
H_at_f = np.full(F_shape, 0.0, dtype)
152+
H_at_f = np.full(F_shape, 0.0, dtype, device=device)
141153

142154
# auxiliary variables for RK time integration
143-
e1 = np.full(T_shape, 0.0, dtype)
144-
u1 = np.full(U_shape, 0.0, dtype)
145-
v1 = np.full(V_shape, 0.0, dtype)
146-
e2 = np.full(T_shape, 0.0, dtype)
147-
u2 = np.full(U_shape, 0.0, dtype)
148-
v2 = np.full(V_shape, 0.0, dtype)
155+
e1 = np.full(T_shape, 0.0, dtype, device=device)
156+
u1 = np.full(U_shape, 0.0, dtype, device=device)
157+
v1 = np.full(V_shape, 0.0, dtype, device=device)
158+
e2 = np.full(T_shape, 0.0, dtype, device=device)
159+
u2 = np.full(U_shape, 0.0, dtype, device=device)
160+
v2 = np.full(V_shape, 0.0, dtype, device=device)
149161

150162
def exact_solution(t, x_t_2d, y_t_2d, x_u_2d, y_u_2d, x_v_2d, y_v_2d):
151163
"""
@@ -174,7 +186,7 @@ def bathymetry(x_t_2d, y_t_2d, lx, ly):
174186
Water depth at rest
175187
"""
176188
bath = 1.0
177-
return bath * np.full(T_shape, 1.0, dtype)
189+
return bath * np.full(T_shape, 1.0, dtype, device=device)
178190

179191
# inital elevation
180192
u0, v0, e0 = exact_solution(0, x_t_2d, y_t_2d, x_u_2d, y_u_2d, x_v_2d, y_v_2d)

‎examples/stencil-2d.py‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -47,15 +47,15 @@
4747
# The output consists of diagnostics to make sure the
4848
# algorithm worked, and of timing statistics.
4949
#
50-
# HISTORY: - Written by Rob Van der Wijngnumpyrt, February 2009.
50+
# HISTORY: - Written by Rob Van der Wijngaart, February 2009.
5151
# - RvdW: Removed unrolling pragmas for clarity;
5252
# added constant to array "in" at end of each iteration to force
5353
# refreshing of neighbor data in parallel versions; August 2013
5454
# - Converted to Python by Jeff Hammond, February 2016.
5555
#
5656
# *******************************************************************
5757

58-
import sys
58+
import sys, os
5959

6060
print(
6161
"Python version = ", str(sys.version_info.major) + "." + str(sys.version_info.minor)
@@ -117,11 +117,13 @@ def main():
117117
print("Data type = double precision")
118118
print("Compact representation of stencil loop body")
119119

120+
device = os.getenv("SHARPY_USE_GPU", "")
121+
120122
# there is certainly a more Pythonic way to initialize W,
121123
# but it will have no impact on performance.
122124
t0 = timer()
123-
W = np.zeros(((2 * r + 1), (2 * r + 1)), dtype=np.float32)
124-
B = np.zeros((n, n), dtype=np.float32)
125+
W = np.zeros(((2 * r + 1), (2 * r + 1)), dtype=np.float32, device=device)
126+
B = np.zeros((n, n), dtype=np.float32, device=device)
125127

126128
if pattern == "star":
127129
stencil_size = 4 * r + 1
@@ -143,7 +145,9 @@ def main():
143145
W[r + j, r + j] = +1.0 / (4 * j * r)
144146
W[r - j, r - j] = -1.0 / (4 * j * r)
145147

146-
A = np.numpy.fromfunction(lambda i, j: i + j, (n, n), dtype=np.float32)
148+
A = np.numpy.fromfunction(
149+
lambda i, j: i + j, (n, n), dtype=np.float32, device=device
150+
)
147151

148152
for k in range(iterations + 1):
149153
# start timer after a warmup iteration

‎examples/wave_equation.py‎

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,10 @@
2828
import time as time_mod
2929
import argparse
3030

31+
import os
32+
33+
device = os.getenv("SHARPY_USE_GPU", "")
34+
3135

3236
def run(n, backend, datatype, benchmark_mode):
3337
if backend == "sharpy":
@@ -94,8 +98,12 @@ def info(s):
9498
t_end = 1.0
9599

96100
# coordinate arrays
97-
x_t_2d = fromfunction(lambda i, j: xmin + i * dx + dx / 2, (nx, ny), dtype=dtype)
98-
y_t_2d = fromfunction(lambda i, j: ymin + j * dy + dy / 2, (nx, ny), dtype=dtype)
101+
x_t_2d = fromfunction(
102+
lambda i, j: xmin + i * dx + dx / 2, (nx, ny), dtype=dtype, device=device
103+
)
104+
y_t_2d = fromfunction(
105+
lambda i, j: ymin + j * dy + dy / 2, (nx, ny), dtype=dtype, device=device
106+
)
99107

100108
T_shape = (nx, ny)
101109
U_shape = (nx + 1, ny)
@@ -111,17 +119,17 @@ def info(s):
111119
info(f"Total DOFs: {dofs_T + dofs_U + dofs_V}")
112120

113121
# prognostic variables: elevation, (u, v) velocity
114-
e = np.full(T_shape, 0.0, dtype)
115-
u = np.full(U_shape, 0.0, dtype)
116-
v = np.full(V_shape, 0.0, dtype)
122+
e = np.full(T_shape, 0.0, dtype, device=device)
123+
u = np.full(U_shape, 0.0, dtype, device=device)
124+
v = np.full(V_shape, 0.0, dtype, device=device)
117125

118126
# auxiliary variables for RK time integration
119-
e1 = np.full(T_shape, 0.0, dtype)
120-
u1 = np.full(U_shape, 0.0, dtype)
121-
v1 = np.full(V_shape, 0.0, dtype)
122-
e2 = np.full(T_shape, 0.0, dtype)
123-
u2 = np.full(U_shape, 0.0, dtype)
124-
v2 = np.full(V_shape, 0.0, dtype)
127+
e1 = np.full(T_shape, 0.0, dtype, device=device)
128+
u1 = np.full(U_shape, 0.0, dtype, device=device)
129+
v1 = np.full(V_shape, 0.0, dtype, device=device)
130+
e2 = np.full(T_shape, 0.0, dtype, device=device)
131+
u2 = np.full(U_shape, 0.0, dtype, device=device)
132+
v2 = np.full(V_shape, 0.0, dtype, device=device)
125133

126134
def exact_elev(t, x_t_2d, y_t_2d, lx, ly):
127135
"""

‎imex_version.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
630ebe45f32b870cf524d05c57dbbe60799c9639
1+
eb68f1d2dcc6ef17099c1f7e8c439a878cb75557

‎src/Creator.cpp‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -79,14 +79,13 @@ struct DeferredFull : public Deferred {
7979
intptr_t o_offset, const intptr_t *o_sizes,
8080
const intptr_t *o_strides, void *r_allocated, void *r_aligned,
8181
intptr_t r_offset, const intptr_t *r_sizes,
82-
const intptr_t *r_strides, uint64_t *lo_allocated,
83-
uint64_t *lo_aligned) {
82+
const intptr_t *r_strides, std::vector<int64_t> &&loffs) {
8483
assert(rank == this->rank());
8584
this->set_value(std::move(mk_tnsr(
8685
this->guid(), _dtype, this->shape(), this->device(), this->team(),
8786
l_allocated, l_aligned, l_offset, l_sizes, l_strides, o_allocated,
8887
o_aligned, o_offset, o_sizes, o_strides, r_allocated, r_aligned,
89-
r_offset, r_sizes, r_strides, lo_allocated, lo_aligned)));
88+
r_offset, r_sizes, r_strides, std::move(loffs))));
9089
});
9190
return false;
9291
}
@@ -141,15 +140,14 @@ struct DeferredArange : public Deferred {
141140
intptr_t o_offset, const intptr_t *o_sizes,
142141
const intptr_t *o_strides, void *r_allocated, void *r_aligned,
143142
intptr_t r_offset, const intptr_t *r_sizes,
144-
const intptr_t *r_strides, uint64_t *lo_allocated,
145-
uint64_t *lo_aligned) {
143+
const intptr_t *r_strides, std::vector<int64_t> &&loffs) {
146144
assert(rank == 1);
147145
assert(o_strides[0] == 1);
148146
this->set_value(std::move(mk_tnsr(
149147
this->guid(), _dtype, this->shape(), this->device(), this->team(),
150148
l_allocated, l_aligned, l_offset, l_sizes, l_strides, o_allocated,
151149
o_aligned, o_offset, o_sizes, o_strides, r_allocated, r_aligned,
152-
r_offset, r_sizes, r_strides, lo_allocated, lo_aligned)));
150+
r_offset, r_sizes, r_strides, std::move(loffs))));
153151
});
154152
return false;
155153
}
@@ -202,15 +200,14 @@ struct DeferredLinspace : public Deferred {
202200
intptr_t o_offset, const intptr_t *o_sizes,
203201
const intptr_t *o_strides, void *r_allocated, void *r_aligned,
204202
intptr_t r_offset, const intptr_t *r_sizes,
205-
const intptr_t *r_strides, uint64_t *lo_allocated,
206-
uint64_t *lo_aligned) {
203+
const intptr_t *r_strides, std::vector<int64_t> &&loffs) {
207204
assert(rank == 1);
208205
assert(l_strides[0] == 1);
209206
this->set_value(std::move(mk_tnsr(
210207
this->guid(), _dtype, this->shape(), this->device(), this->team(),
211208
l_allocated, l_aligned, l_offset, l_sizes, l_strides, o_allocated,
212209
o_aligned, o_offset, o_sizes, o_strides, r_allocated, r_aligned,
213-
r_offset, r_sizes, r_strides, lo_allocated, lo_aligned)));
210+
r_offset, r_sizes, r_strides, std::move(loffs))));
214211
});
215212
return false;
216213
}

‎src/Deferred.cpp‎

Lines changed: 4 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,6 @@
1515
#include "include/sharpy/itac.hpp"
1616
#include "include/sharpy/jit/mlir.hpp"
1717

18-
#include <imex/Dialect/NDArray/IR/NDArrayOps.h>
19-
#include <mlir/Dialect/Func/IR/FuncOps.h>
20-
#include <mlir/Dialect/LLVMIR/LLVMDialect.h>
2118
#include <oneapi/tbb/concurrent_queue.h>
2219

2320
#include <pybind11/pybind11.h>
@@ -96,36 +93,13 @@ void process_promises() {
9693
std::vector<Runable::ptr_type> deleters;
9794

9895
do {
99-
::mlir::OpBuilder builder(&jit.context());
100-
auto loc = builder.getUnknownLoc();
101-
102-
// Create a MLIR module
103-
auto module = builder.create<::mlir::ModuleOp>(loc);
104-
// Create the jit func
105-
// create dummy type, we'll replace it with the actual type later
106-
auto dummyFuncType = builder.getFunctionType({}, {});
107-
if (false) {
108-
::mlir::OpBuilder::InsertionGuard guard(builder);
109-
// Insert before module terminator.
110-
builder.setInsertionPoint(module.getBody(),
111-
std::prev(module.getBody()->end()));
112-
auto func = builder.create<::mlir::func::FuncOp>(loc, "_debugFunc",
113-
dummyFuncType);
114-
func.setPrivate();
115-
}
116-
std::string fname("sharpy_jit");
117-
auto function =
118-
builder.create<::mlir::func::FuncOp>(loc, fname, dummyFuncType);
119-
// create function entry block
120-
auto &entryBlock = *function.addEntryBlock();
121-
// Set the insertion point in the builder to the beginning of the function
122-
// body
123-
builder.setInsertionPointToStart(&entryBlock);
12496
// we need to keep runables/deferred/futures alive until we set their values
12597
// below
12698
std::vector<Runable::ptr_type> runables;
12799

128-
jit::DepManager dm(function);
100+
jit::DepManager dm(jit);
101+
auto &builder = dm.getBuilder();
102+
auto loc = builder.getUnknownLoc();
129103
Runable::ptr_type d;
130104

131105
if (!deleters.empty()) {
@@ -160,30 +134,7 @@ void process_promises() {
160134
}
161135

162136
if (!runables.empty()) {
163-
// get input buffers (before results!)
164-
auto input = std::move(dm.finalize_inputs());
165-
// create return statement and adjust function type
166-
uint64_t osz = dm.handleResult(builder);
167-
// also request generation of c-wrapper function
168-
function->setAttr(::mlir::LLVM::LLVMDialect::getEmitCWrapperAttrName(),
169-
builder.getUnitAttr());
170-
if (jit.verbose())
171-
function.getFunctionType().dump();
172-
// add the function to the module
173-
module.push_back(function);
174-
175-
if (osz > 0 || !input.empty()) {
176-
// compile and run the module
177-
auto output = jit.run(module, fname, input, osz);
178-
if (output.size() != osz)
179-
throw std::runtime_error("failed running jit");
180-
181-
// push results to deliver promises
182-
dm.deliver(output, osz);
183-
} else {
184-
if (jit.verbose())
185-
std::cerr << "\tskipping\n";
186-
}
137+
dm.finalizeAndRun();
187138
} // no else needed
188139

189140
// now we execute the deferred action which could not be compiled

‎src/EWBinOp.cpp‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -145,13 +145,12 @@ struct DeferredEWBinOp : public Deferred {
145145
intptr_t o_offset, const intptr_t *o_sizes,
146146
const intptr_t *o_strides, void *r_allocated, void *r_aligned,
147147
intptr_t r_offset, const intptr_t *r_sizes,
148-
const intptr_t *r_strides, uint64_t *lo_allocated,
149-
uint64_t *lo_aligned) {
148+
const intptr_t *r_strides, std::vector<int64_t> &&loffs) {
150149
this->set_value(std::move(mk_tnsr(
151150
this->guid(), _dtype, this->shape(), this->device(), this->team(),
152151
l_allocated, l_aligned, l_offset, l_sizes, l_strides, o_allocated,
153152
o_aligned, o_offset, o_sizes, o_strides, r_allocated, r_aligned,
154-
r_offset, r_sizes, r_strides, lo_allocated, lo_aligned)));
153+
r_offset, r_sizes, r_strides, std::move(loffs))));
155154
});
156155
return false;
157156
}

‎src/EWUnyOp.cpp‎

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -221,26 +221,26 @@ struct DeferredEWUnyOp : public Deferred {
221221
// positive op will be eliminated so it is equivalent to a view
222222
auto view = ndOpId == ::imex::ndarray::POSITIVE;
223223

224-
dm.addVal(
225-
this->guid(), uop,
226-
[this, view](uint64_t rank, void *l_allocated, void *l_aligned,
227-
intptr_t l_offset, const intptr_t *l_sizes,
228-
const intptr_t *l_strides, void *o_allocated,
229-
void *o_aligned, intptr_t o_offset,
230-
const intptr_t *o_sizes, const intptr_t *o_strides,
231-
void *r_allocated, void *r_aligned, intptr_t r_offset,
232-
const intptr_t *r_sizes, const intptr_t *r_strides,
233-
uint64_t *lo_allocated, uint64_t *lo_aligned) {
234-
auto t = mk_tnsr(this->guid(), _dtype, this->shape(), this->device(),
235-
this->team(), l_allocated, l_aligned, l_offset,
236-
l_sizes, l_strides, o_allocated, o_aligned, o_offset,
237-
o_sizes, o_strides, r_allocated, r_aligned, r_offset,
238-
r_sizes, r_strides, lo_allocated, lo_aligned);
239-
if (view && Registry::has(_a)) {
240-
t->set_base(Registry::get(_a).get());
241-
}
242-
this->set_value(std::move(t));
243-
});
224+
dm.addVal(this->guid(), uop,
225+
[this, view](
226+
uint64_t rank, void *l_allocated, void *l_aligned,
227+
intptr_t l_offset, const intptr_t *l_sizes,
228+
const intptr_t *l_strides, void *o_allocated, void *o_aligned,
229+
intptr_t o_offset, const intptr_t *o_sizes,
230+
const intptr_t *o_strides, void *r_allocated, void *r_aligned,
231+
intptr_t r_offset, const intptr_t *r_sizes,
232+
const intptr_t *r_strides, std::vector<int64_t> &&loffs) {
233+
auto t = mk_tnsr(this->guid(), _dtype, this->shape(),
234+
this->device(), this->team(), l_allocated,
235+
l_aligned, l_offset, l_sizes, l_strides,
236+
o_allocated, o_aligned, o_offset, o_sizes,
237+
o_strides, r_allocated, r_aligned, r_offset,
238+
r_sizes, r_strides, std::move(loffs));
239+
if (view && Registry::has(_a)) {
240+
t->set_base(Registry::get(_a).get());
241+
}
242+
this->set_value(std::move(t));
243+
});
244244
return false;
245245
}
246246

0 commit comments

Comments
 (0)