Skip to content

Commit

Permalink
[CIR] Simplify channel syntax and add range interface to buffer
Browse files Browse the repository at this point in the history 
Improve also documentation and type-safety.
  • Loading branch information
@keryell
keryell committed Dec 6, 2024
1 parent 171b5b2 commit 3b582cd
Show file tree
Hide file tree
Showing 2 changed files with 106 additions and 64 deletions.
162 changes: 102 additions & 60 deletions include/aie/CIR/runtime/aie++/aie++.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,59 +16,10 @@ template <typename T, std::size_t Size> struct buffer {
operator storage_t&() { return storage; }
// The previous is not enough
decltype(auto) operator[](std::size_t index) { return storage[index]; }
};

// Typically compiled as:
// !ty_22aie3A3Atile3C12C_43E22 = !cir.struct<struct "aie::tile<1, 4>"
// {!cir.int<u, 8>}>
//
// The tile depends on a Device since we can have programs with different
// devices at the same time
template <int X, int Y, typename Device> struct tile {
static constexpr auto x() { return X; }
static constexpr auto y() { return Y; }
// Only tiles from a same device are comparable
template <int X2, int Y2>
friend constexpr auto operator==(const tile&, const tile<X2, Y2, Device>&) {
return X == X2 && Y == Y2;
};
/* template <int X1, int Y1, int X2, int Y2>
friend auto operator<=>(const tile_t<X1, Y1> &, const tile_t<X2, Y2> &) {
return std::strong_ordering::equal;
};*/
/*
friend constexpr auto operator<=>(const tile &a, const tile &b) {
// return operator<=>(std::array { a.x, a.y }, std::array { b.x, b.y });
return operator<=>(a.x, a.y);
}
*/

template <typename Code> static inline Code tile_code;
auto begin() { return storage.begin(); }

template <typename T, std::size_t Size>
buffer<T, Size> buffer() __attribute__((annotate("aie.tile.buffer"))) {
return {};
}

// Typically compiled as:
// cir.call @_ZN3aie4tileILi1ELi4EE7programIZ4mainE3$_0EEvOT_(%2, %7) :
// (!cir.ptr<!ty_22aie3A3Atile3C12C_43E22>, !cir.ptr<!ty_22anon2E122>) -> ()
// loc(#loc63)
void program(auto&& code) __attribute__((annotate("aie.tile.program"))) {
// Use this since std::function crashes ClangIR 2024/09/12
// tile_code<decltype(&code)> = &code;
// Just to instantiate the lambda body while waiting for std::function
code();
}
};

// template <int X, int Y> inline constexpr tile<X, Y> tile;

template <typename Storage> struct tile_handle {
Storage tile_memory;
constexpr tile_handle(Storage tile_memory)
: tile_memory { tile_memory } {};
constexpr auto& mem() { return tile_memory; }
auto end() { return storage.end(); }
};

template <typename Channel> void aquire() {}
Expand Down Expand Up @@ -102,9 +53,18 @@ template <typename Channel> struct accessor {
~accessor() { release<Channel>(); }
};

// Channel abstraction used to send data between tiles.
//
// This is lowered to MLIR aie.objectfifo.
template <typename ValueType, std::size_t Size, typename FromTile,
typename ToTile>
struct channel {
static_assert(std::is_same_v<typename FromTile::device_type,
typename ToTile::device_type>,
"Only tiles from the same device can be connected");
using from_tile_type = FromTile;
using to_tile_type = ToTile;
using device_type = typename to_tile_type::device_type;
// clang++:
// /home/rkeryell/Xilinx/Projects/LLVM/worktrees/clangir/clang/lib/CIR/CodeGen/CIRGenExpr.cpp:552:
// CIRGenCallee cir::CIRGenFunction::buildCallee(const clang::Expr *):
Expand Down Expand Up @@ -134,9 +94,97 @@ struct channel {
}
};

// Represent the tile in AIE Device.
//
// X is the column number starting at 0.
//
// Y is the row number starting at 0.
//
// Typically compiled as:
// !ty_22aie3A3Atile3C12C_43E22 = !cir.struct<struct "aie::tile<1, 4>"
// {!cir.int<u, 8>}>
//
// The tile depends on a Device since we can have programs with different
// devices at the same time
template <int X, int Y, typename Device> struct tile {
using device_type = Device;
static constexpr auto x() { return X; }
static constexpr auto y() { return Y; }
// Only tiles from a same device are comparable
template <int X2, int Y2>
friend constexpr auto operator==(const tile&, const tile<X2, Y2, Device>&) {
return X == X2 && Y == Y2;
};
/* template <int X1, int Y1, int X2, int Y2>
friend auto operator<=>(const tile_t<X1, Y1> &, const tile_t<X2, Y2> &) {
return std::strong_ordering::equal;
};*/
/*
friend constexpr auto operator<=>(const tile &a, const tile &b) {
// return operator<=>(std::array { a.x, a.y }, std::array { b.x, b.y });
return operator<=>(a.x, a.y);
}
*/

template <typename Code> static inline Code tile_code;

// Get a buffer with Size elements of type T.
template <typename T, std::size_t Size>
buffer<T, Size> buffer() __attribute__((annotate("aie.tile.buffer"))) {
return {};
}

// Define the code to be run inside a tile.
//
// Typically compiled as:
// cir.call @_ZN3aie4tileILi1ELi4EE7programIZ4mainE3$_0EEvOT_(%2, %7) :
// (!cir.ptr<!ty_22aie3A3Atile3C12C_43E22>, !cir.ptr<!ty_22anon2E122>) -> ()
// loc(#loc63)
void program(auto&& code) __attribute__((annotate("aie.tile.program"))) {
// Use this since std::function crashes ClangIR 2024/09/12
// tile_code<decltype(&code)> = &code;
// Just to instantiate the lambda body while waiting for std::function
code();
}

// Create a channel to communicate between 2 tiles.
template <typename ValueType, std::size_t Size, typename ToTile>
channel<ValueType, Size, tile, ToTile> channel_to(ToTile& to_tile,
std::size_t capacity)
__attribute__((annotate("aie.device.channel"))) {
return { *this, to_tile, capacity };
}
};

// template <int X, int Y> inline constexpr tile<X, Y> tile;

template <typename Storage> struct tile_handle {
Storage tile_memory;
constexpr tile_handle(Storage tile_memory)
: tile_memory { tile_memory } {};
constexpr auto& mem() { return tile_memory; }
};

/*
template <typename ValueType, std::size_t Size, typename FromTile,
typename ToTile>
channel<ValueType, Size>(FromTile& from_tile, ToTile& to_tile,
std::size_t capacity)
->channel<ValueType, Size, FromTile, ToTile>;
*/
// Inject in aie:: to ease use
enum : std::int8_t { npu1 = 42 };
enum : std::int8_t {
npu1 = 42, //< AIE2 device found in RyzenAI 9 7940HS
npu2,
npu3,
npu4,
pacifica //< Fancy AIE device found on Californian shore
};

// Abstraction representing an AIE device in the system.
//
// DeviceModel specifies an architecture model, such as aie::npu1.
//
// Typically compiled as:
// !ty_aie3A3Adevice3Caie3A3Anpu12C_aie3A3A28lambda_at_2E2Faie2B2B2Ehpp3A763A54293E
// = !cir.struct<struct "aie::device<aie::npu1, aie::(lambda at
Expand All @@ -147,21 +195,15 @@ enum : std::int8_t { npu1 = 42 };
// type
template <auto DeviceModel = npu1, typename Unique = decltype([] {})>
struct device {
// Get the tile of a device. X is the column of the tile, Y is the row of the
// tile.
template <int X, int Y>
tile<X, Y, device> tile()
__attribute__((annotate("aie.device.tile", X, Y, DeviceModel,
std::to_underlying(DeviceModel)))) {
return {};
}

template <typename ValueType, std::size_t Size, typename FromTile,
typename ToTile>
channel<ValueType, Size, FromTile, ToTile>
channel(FromTile& from_tile, ToTile& to_tile, std::size_t capacity)
__attribute__((annotate("aie.device.channel"))) {
return { from_tile, to_tile, capacity };
}

void constexpr run() {}
};

Expand Down
8 changes: 4 additions & 4 deletions test/CIR/aie++/channel.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
Iron MICRO24 slide 34..
Iron MICRO24 slide 34.
https://github.com/Xilinx/mlir-aie/blob/main/mlir_tutorials
*/

Expand All @@ -18,15 +18,15 @@ int main() {
aie::device<aie::npu1> d;
auto a = d.tile<1, 3>();
auto b = d.tile<2, 3>();
auto of0 = d.channel<std::int32_t, 256>(a, b, 3);
auto of = a.channel_to<std::int32_t, 256>(b, 3);
a.program([&] {
for (int i = 0; i < 3; ++i) {
auto acc = of0.out_acquire_release(1);
auto acc = of.out_acquire_release(1);
produce(acc);
}
});
b.program([&] {
auto acc = of0.in_acquire_release(3);
auto acc = of.in_acquire_release(3);
consume(acc[0]);
consume(acc[1]);
consume(acc[2]);
Expand Down

0 comments on commit 3b582cd

Please sign in to comment.