Skip to content
This repository was archived by the owner on Jan 26, 2026. It is now read-only.

Commit 101a8b0

Browse files
fschlimbfschlimb
authored
adding ITAC instrumentation (#53)
* adding ITAC instrumentation
1 parent a4a0019 commit 101a8b0

File tree

4 files changed

+101
-3
lines changed

4 files changed

+101
-3
lines changed

‎src/Deferred.cpp‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
#include "include/ddptensor/Registry.hpp"
1313
#include "include/ddptensor/Service.hpp"
1414
#include "include/ddptensor/Transceiver.hpp"
15+
#include "include/ddptensor/itac.hpp"
1516

1617
#include <imex/Dialect/PTensor/IR/PTensorOps.h>
1718
#include <mlir/Dialect/Func/IR/FuncOps.h>
@@ -85,8 +86,15 @@ void Runable::fini() { _deferred.clear(); }
8586
// statement) is finalized, the function gets compiled and executed. The loop
8687
// completes by calling run() on the requesting object.
8788
void process_promises() {
89+
int vtProcessSym, vtDDPTClass, vtPopSym;
90+
VT(VT_classdef, "ddpt", &vtDDPTClass);
91+
VT(VT_funcdef, "process", vtDDPTClass, &vtProcessSym);
92+
VT(VT_funcdef, "pop", vtDDPTClass, &vtPopSym);
93+
VT(VT_begin, vtProcessSym);
94+
8895
bool done = false;
8996
jit::JIT jit;
97+
9098
do {
9199
::mlir::OpBuilder builder(&jit.context());
92100
auto loc = builder.getUnknownLoc();
@@ -121,7 +129,9 @@ void process_promises() {
121129

122130
Runable::ptr_type d;
123131
while (true) {
132+
VT(VT_begin, vtPopSym);
124133
_deferred.pop(d);
134+
VT(VT_end, vtPopSym);
125135
if (d) {
126136
if (d->generate_mlir(builder, loc, dm)) {
127137
break;

‎src/ddptensor.cpp‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ using namespace pybind11::literals; // to bring _a
4141
#include "ddptensor/Service.hpp"
4242
#include "ddptensor/SetGetItem.hpp"
4343
#include "ddptensor/Sorting.hpp"
44+
#include "ddptensor/itac.hpp"
4445
#include "ddptensor/jit/mlir.hpp"
4546

4647
#include <iostream>
@@ -96,31 +97,49 @@ void init(bool cw) {
9697
}
9798

9899
void sync_promises() {
100+
int vtWaitSym, vtDDPTClass;
101+
VT(VT_classdef, "ddpt", &vtDDPTClass);
102+
VT(VT_funcdef, "wait", vtDDPTClass, &vtWaitSym);
103+
VT(VT_begin, vtWaitSym);
99104
py::gil_scoped_release release;
100105
(void)Service::run().get();
106+
VT(VT_end, vtWaitSym);
101107
}
102108

103109
// #########################################################################
104110

105111
/// trigger compile&run and return future value
106112
#define PY_SYNC_RETURN(_f) \
113+
int vtWaitSym, vtDDPTClass; \
114+
VT(VT_classdef, "ddpt", &vtDDPTClass); \
115+
VT(VT_funcdef, "wait", vtDDPTClass, &vtWaitSym); \
116+
VT(VT_begin, vtWaitSym); \
107117
py::gil_scoped_release release; \
108118
Service::run(); \
109-
return (_f).get()
119+
auto r = (_f).get(); \
120+
VT(VT_end, vtWaitSym); \
121+
return r
110122

111123
/// trigger compile&run and return given attribute _x
112124
#define SYNC_RETURN(_f, _a) \
125+
int vtWaitSym, vtDDPTClass; \
126+
VT(VT_classdef, "ddpt", &vtDDPTClass); \
127+
VT(VT_funcdef, "wait", vtDDPTClass, &vtWaitSym); \
128+
VT(VT_begin, vtWaitSym); \
113129
py::gil_scoped_release release; \
114130
Service::run(); \
115-
return (_f).get().get()->_a()
131+
auto r = (_f).get().get()->_a(); \
132+
VT(VT_end, vtWaitSym); \
133+
return r
116134

117-
/// Rerplicate ddptensor/future and SYNC_RETURN attributre _a
135+
/// Replicate ddptensor/future and SYNC_RETURN attribute _a
118136
#define REPL_SYNC_RETURN(_f, _a) \
119137
auto r_ = std::unique_ptr<ddptensor>(Service::replicate(f)); \
120138
SYNC_RETURN(r_->get(), _a)
121139

122140
// Finally our Python module
123141
PYBIND11_MODULE(_ddptensor, m) {
142+
124143
initFactories();
125144

126145
jit::init();

‎src/include/ddptensor/itac.hpp‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
#pragma once
2+
3+
#ifndef NDEBUG
4+
5+
#define VT_ErrorCode_t int // it is actually an enum...
6+
extern "C" {
7+
VT_ErrorCode_t VT_classdef(const char *, int *) __attribute__((weak));
8+
VT_ErrorCode_t VT_funcdef(const char *, int, int *) __attribute__((weak));
9+
VT_ErrorCode_t VT_begin(int) __attribute__((weak));
10+
VT_ErrorCode_t VT_end(int) __attribute__((weak));
11+
}
12+
#define VT(_sym, ...) \
13+
if (_sym != NULL) { \
14+
_sym(__VA_ARGS__); \
15+
}
16+
#define HAS_ITAC() (VT_classdef != NULL)
17+
18+
#else
19+
20+
#define VT(...) ;
21+
#define HAS_ITAC() false
22+
23+
#endif

‎src/jit/mlir.cpp‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,8 @@
101101
#include "llvm/Support/TargetSelect.h"
102102
// #include "llvm/Support/raw_ostream.h"
103103

104+
#include "ddptensor/itac.hpp"
105+
104106
namespace jit {
105107

106108
static ::mlir::Type makeSignlessType(::mlir::Type type) {
@@ -261,6 +263,22 @@ uint64_t DepManager::handleResult(::mlir::OpBuilder &builder) {
261263
++idx;
262264
}
263265

266+
if (HAS_ITAC()) {
267+
int vtExeSym, vtDDPTClass;
268+
VT(VT_classdef, "ddpt", &vtDDPTClass);
269+
VT(VT_funcdef, "execute", vtDDPTClass, &vtExeSym);
270+
auto s = builder.create<::mlir::arith::ConstantOp>(
271+
loc, builder.getI32IntegerAttr(vtExeSym));
272+
auto end = builder.create<::mlir::func::CallOp>(
273+
builder.getUnknownLoc(), "VT_end",
274+
::mlir::TypeRange(builder.getIntegerType(32)), ::mlir::ValueRange(s));
275+
mlir::OpBuilder::InsertionGuard guard(builder);
276+
builder.setInsertionPointToStart(end->getBlock());
277+
(void)builder.create<::mlir::func::CallOp>(
278+
builder.getUnknownLoc(), "VT_begin",
279+
::mlir::TypeRange(builder.getIntegerType(32)), ::mlir::ValueRange(s));
280+
}
281+
264282
// add return statement
265283
auto ret_value = builder.create<::mlir::func::ReturnOp>(
266284
builder.getUnknownLoc(), ret_values);
@@ -348,28 +366,55 @@ std::vector<intptr_t> JIT::run(::mlir::ModuleOp &module,
348366
const std::string &fname,
349367
std::vector<void *> &inp, size_t osz) {
350368

369+
int vtDDPTClass, vtHashSym, vtEEngineSym, vtRunSym, vtHashGenSym;
370+
if (HAS_ITAC()) {
371+
VT(VT_classdef, "ddpt", &vtDDPTClass);
372+
VT(VT_funcdef, "lookup_cache", vtDDPTClass, &vtHashSym);
373+
VT(VT_funcdef, "gen_sha", vtDDPTClass, &vtHashGenSym);
374+
VT(VT_funcdef, "eengine", vtDDPTClass, &vtEEngineSym);
375+
VT(VT_funcdef, "run", vtDDPTClass, &vtRunSym);
376+
VT(VT_begin, vtEEngineSym);
377+
378+
::mlir::OpBuilder builder(module->getContext());
379+
::mlir::OpBuilder::InsertionGuard guard(builder);
380+
builder.setInsertionPoint(module.getBody(),
381+
std::prev(module.getBody()->end()));
382+
auto intTyp = builder.getIntegerType(32);
383+
auto funcType = builder.getFunctionType({intTyp}, {intTyp});
384+
builder.create<::mlir::func::FuncOp>(module.getLoc(), "VT_begin", funcType)
385+
.setPrivate();
386+
builder.create<::mlir::func::FuncOp>(module.getLoc(), "VT_end", funcType)
387+
.setPrivate();
388+
}
389+
351390
::mlir::ExecutionEngine *enginePtr;
352391
std::unique_ptr<::mlir::ExecutionEngine> tmpEngine;
353392

354393
if (_useCache) {
394+
VT(VT_begin, vtHashGenSym);
355395
static std::map<std::array<unsigned char, 20>,
356396
std::unique_ptr<::mlir::ExecutionEngine>>
357397
engineCache;
358398

359399
llvm::raw_sha1_ostream shaOS;
360400
module->print(shaOS);
361401
auto cksm = shaOS.sha1();
402+
VT(VT_end, vtHashGenSym);
362403

404+
VT(VT_begin, vtHashSym);
363405
if (auto search = engineCache.find(cksm); search == engineCache.end()) {
364406
engineCache[cksm] = createExecutionEngine(module);
365407
} else {
366408
if (_verbose)
367409
std::cerr << "cached..." << std::endl;
368410
}
369411
enginePtr = engineCache[cksm].get();
412+
VT(VT_end, vtHashSym);
370413
} else {
414+
VT(VT_begin, vtHashSym);
371415
tmpEngine = createExecutionEngine(module);
372416
enginePtr = tmpEngine.get();
417+
VT(VT_end, vtHashSym);
373418
}
374419

375420
auto expectedFPtr =
@@ -398,6 +443,7 @@ std::vector<intptr_t> JIT::run(::mlir::ModuleOp &module,
398443
// call function
399444
(*jittedFuncPtr)(args.data());
400445

446+
VT(VT_end, vtEEngineSym);
401447
return out;
402448
}
403449

0 commit comments

Comments
 (0)