Skip to content

Add comm engine attrs for collective and async ops #723

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
FangRui0 wants to merge 1 commit into
base: main
Choose a base branch
from
+367 −116
Open

Add comm engine attrs for collective and async ops #723

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 18 additions & 7 deletions docs/PTO_IR_manual.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9192,6 +9192,7 @@ This section documents PTO communication primitives. PTOAS currently exposes:
|------|------|-------------|
| `scratch` | `pto.tile_buf` / local memref | Local scratch/staging buffer used by the async runtime |
| `workspace` | `!pto.ptr<...>` / GM memref | Global workspace backing the async session |
| `dmaEngine` | optional `#pto<dma_engine ...>` attr | DMA engine template parameter, defaults to `#pto<dma_engine sdma>` |
| `sync_id` | optional `i32` attr | Session synchronization ID |
| `block_bytes` | optional `i64` attr | Communication block size in bytes |
| `comm_block_offset` | optional `i64` attr | Per-block GM offset in bytes |
Expand All @@ -9204,12 +9205,13 @@ This section documents PTO communication primitives. PTOAS currently exposes:

- `scratch` must be tile-like local storage.
- `workspace` must be a GM pointer/memref.
- `dmaEngine` currently lowers as a compile-time engine selector; omitted means `SDMA`.
- Optional attrs are forwarded as session configuration and must use the declared integer types.

**Basic Example:**

```mlir
%session = pto.comm.build_async_session(%scratch, %workspace : !pto.tile_buf<loc=vec, dtype=i8, rows=1, cols=256, v_row=1, v_col=256, blayout=row_major, slayout=none_box, fractal=512, pad=0>, !pto.ptr<i8>) {sync_id = 0 : i32} -> !pto.async_session
%session = pto.comm.build_async_session(%scratch, %workspace : !pto.tile_buf<loc=vec, dtype=i8, rows=1, cols=256, v_row=1, v_col=256, blayout=row_major, slayout=none_box, fractal=512, pad=0>, !pto.ptr<i8>) {dmaEngine = #pto<dma_engine urma>, sync_id = 0 : i32} -> !pto.async_session
```

---
Expand All @@ -9225,6 +9227,7 @@ This section documents PTO communication primitives. PTOAS currently exposes:
| `dst` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Remote destination buffer |
| `src` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Local source buffer |
| `session` | `!pto.async_session` | Async DMA session |
| `dmaEngine` | optional `#pto<dma_engine ...>` attr | Transfer engine template parameter, defaults to `#pto<dma_engine sdma>` |

**Results:** `!pto.async_event`

Expand All @@ -9233,11 +9236,12 @@ This section documents PTO communication primitives. PTOAS currently exposes:
- `dst` / `src` must be GM-shaped values with identical element type and static shape.
- Current lowering only supports flat contiguous logical-1D transfers for async GM operands.
- `session` must come from `pto.comm.build_async_session`.
- `dmaEngine` should match the engine chosen for the session when non-default lowering is required.

**Basic Example:**

```mlir
%event = pto.comm.tput_async(%dst, %src, %session : !pto.partition_tensor_view<128xf32>, !pto.partition_tensor_view<128xf32>, !pto.async_session) -> !pto.async_event
%event = pto.comm.tput_async(%dst, %src, %session : !pto.partition_tensor_view<128xf32>, !pto.partition_tensor_view<128xf32>, !pto.async_session) {dmaEngine = #pto<dma_engine urma>} -> !pto.async_event
```

---
Expand All @@ -9253,18 +9257,20 @@ This section documents PTO communication primitives. PTOAS currently exposes:
| `dst` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Local destination buffer |
| `src` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Remote source buffer |
| `session` | `!pto.async_session` | Async DMA session |
| `dmaEngine` | optional `#pto<dma_engine ...>` attr | Transfer engine template parameter, defaults to `#pto<dma_engine sdma>` |

**Results:** `!pto.async_event`

**Constraints & Verification:**

- Same operand constraints as `pto.comm.tput_async`.
- `session` must be compatible with the transfer workspace and staging configuration.
- `dmaEngine` is forwarded to emitc as the compile-time `pto::comm::DmaEngine` template argument.

**Basic Example:**

```mlir
%event = pto.comm.tget_async(%dst, %src, %session : !pto.partition_tensor_view<128xf32>, !pto.partition_tensor_view<128xf32>, !pto.async_session) -> !pto.async_event
%event = pto.comm.tget_async(%dst, %src, %session : !pto.partition_tensor_view<128xf32>, !pto.partition_tensor_view<128xf32>, !pto.async_session) {dmaEngine = #pto<dma_engine urma>} -> !pto.async_event
```

---
Expand Down Expand Up @@ -9393,13 +9399,15 @@ pto.comm.twait(%sig, %v : !pto.partition_tensor_view<1xi32>, i32) {cmp = #pto<wa
| `src` | GM-shaped value | Root source buffer |
| `recv` | `recv(%ping)` or `recv(%ping, %pong)` | One or two local VEC staging tiles |
| `group` | variadic GM-shaped values | Parallel group members |
| `collEngine` | optional `#pto<coll_engine ...>` attr | Collective engine template parameter, defaults to `#pto<coll_engine aiv>` |
| `root` | `i32` attr | Root rank index inside `group` |

**Constraints & Verification:**

- `group` must be non-empty and all members must have identical types.
- `src` must have the same type as each `group` member.
- `root` must be in range `[0, group.size)`.
- `collEngine` is forwarded to emitc as the compile-time `pto::comm::CollEngine` template argument.

**Examples:**

Expand All @@ -9411,7 +9419,7 @@ pto.comm.tbroadcast(%src, recv(%ping), group(%g0, %g1, %g2) :
!pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=128, v_row=1, v_col=128, blayout=row_major, slayout=none_box, fractal=512, pad=0>,
!pto.partition_tensor_view<128xf32>,
!pto.partition_tensor_view<128xf32>,
!pto.partition_tensor_view<128xf32>) {root = 1 : i32}
!pto.partition_tensor_view<128xf32>) {collEngine = #pto<coll_engine ccu>, root = 1 : i32}
```

Optional ping–pong (`recv(%ping, %pong)` adds a second tile type in the operand-type list):
Expand Down Expand Up @@ -9439,6 +9447,7 @@ pto.comm.tbroadcast(%src, recv(%ping, %pong), group(%g0, %g1, %g2) :
| `dst` | GM-shaped value | Destination buffer (gather target) |
| `recv` | `recv(%ping)` or `recv(%ping, %pong)` | Staging tile(s) |
| `group` | variadic GM-shaped values | Parallel group members |
| `collEngine` | optional `#pto<coll_engine ...>` attr | Collective engine template parameter, defaults to `#pto<coll_engine aiv>` |
| `root` | `i32` attr | Root rank index inside `group` |

**Constraints & Verification:**
Expand All @@ -9463,7 +9472,7 @@ pto.comm.tgather(%dst, recv(%ping, %pong), group(%g0, %g1, %g2) :
!pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=128, v_row=1, v_col=128, blayout=row_major, slayout=none_box, fractal=512, pad=0>,
!pto.partition_tensor_view<128xf32>,
!pto.partition_tensor_view<128xf32>,
!pto.partition_tensor_view<128xf32>) {root = 1 : i32}
!pto.partition_tensor_view<128xf32>) {collEngine = #pto<coll_engine ccu>, root = 1 : i32}
```

---
Expand All @@ -9479,6 +9488,7 @@ pto.comm.tgather(%dst, recv(%ping, %pong), group(%g0, %g1, %g2) :
| `src` | GM-shaped value | Source buffer (scatter root) |
| `recv` | `recv(%ping)` or `recv(%ping, %pong)` | Staging tile(s) |
| `group` | variadic GM-shaped values | Parallel group members |
| `collEngine` | optional `#pto<coll_engine ...>` attr | Collective engine template parameter, defaults to `#pto<coll_engine aiv>` |
| `root` | `i32` attr | Root rank index inside `group` |

**Constraints & Verification:**
Expand All @@ -9503,7 +9513,7 @@ pto.comm.tscatter(%src, recv(%ping, %pong), group(%g0, %g1, %g2) :
!pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=128, v_row=1, v_col=128, blayout=row_major, slayout=none_box, fractal=512, pad=0>,
!pto.partition_tensor_view<128xf32>,
!pto.partition_tensor_view<128xf32>,
!pto.partition_tensor_view<128xf32>) {root = 1 : i32}
!pto.partition_tensor_view<128xf32>) {collEngine = #pto<coll_engine ccu>, root = 1 : i32}
```

---
Expand All @@ -9520,6 +9530,7 @@ pto.comm.tscatter(%src, recv(%ping, %pong), group(%g0, %g1, %g2) :
| `acc` | local VEC tile-like value | Accumulation tile |
| `recv` | `recv(%ping)` or `recv(%ping, %pong)` | One or two receive staging tiles |
| `group` | variadic GM-shaped values | Parallel group members |
| `collEngine` | optional `#pto<coll_engine ...>` attr | Collective engine template parameter, defaults to `#pto<coll_engine aiv>` |
| `reduceOp` | `#pto<reduce_op sum>` / `#pto<reduce_op max>` / `#pto<reduce_op min>` | Reduction mode |
| `root` | `i32` attr | Root rank index inside `group` |

Expand Down Expand Up @@ -9553,7 +9564,7 @@ pto.comm.treduce(%dst, %acc, recv(%ping, %pong), group(%g0, %g1, %g2) :
!pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=128, v_row=1, v_col=128, blayout=row_major, slayout=none_box, fractal=512, pad=0>,
!pto.partition_tensor_view<128xf32>,
!pto.partition_tensor_view<128xf32>,
!pto.partition_tensor_view<128xf32>) {reduceOp = #pto<reduce_op max>, root = 1 : i32}
!pto.partition_tensor_view<128xf32>) {collEngine = #pto<coll_engine ccu>, reduceOp = #pto<reduce_op max>, root = 1 : i32}
```

---
Expand Down
20 changes: 20 additions & 0 deletions include/PTO/IR/PTOAttrs.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -505,6 +505,26 @@ def PTO_ReduceOpAttr : EnumAttr<PTO_Dialect, PTO_ReduceOpEnum, "reduce_op"> {
let summary = "communication reduce operation attribute";
}

def PTO_DmaEngineEnum : PTO_I32Enum<
"DmaEngine", "PTO communication DMA engine", [
I32EnumAttrCase<"SDMA", 0, "sdma">,
I32EnumAttrCase<"URMA", 1, "urma">
]>;

def PTO_DmaEngineAttr : EnumAttr<PTO_Dialect, PTO_DmaEngineEnum, "dma_engine"> {
let summary = "communication DMA engine attribute";
}

def PTO_CollEngineEnum : PTO_I32Enum<
"CollEngine", "PTO communication collective engine", [
I32EnumAttrCase<"AIV", 0, "aiv">,
I32EnumAttrCase<"CCU", 1, "ccu">
]>;

def PTO_CollEngineAttr : EnumAttr<PTO_Dialect, PTO_CollEngineEnum, "coll_engine"> {
let summary = "communication collective engine attribute";
}

def PTO_ReluPreModeEnum : PTO_I32Enum<
"ReluPreMode", "PTO TSTORE relu pre mode", [
I32EnumAttrCase<"NoRelu", 0, "no_relu">,
Expand Down
11 changes: 9 additions & 2 deletions include/PTO/IR/PTOOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1737,6 +1737,7 @@ def BuildAsyncSessionOp : PTO_Op<"comm.build_async_session", [
let arguments = (ins
TileBufOrMemRef:$scratch,
ScalarPtrOrMemRef:$workspace,
DefaultValuedAttr<PTO_DmaEngineAttr, "::mlir::pto::DmaEngine::SDMA">:$dmaEngine,
OptionalAttr<I32Attr>:$sync_id,
OptionalAttr<I64Attr>:$block_bytes,
OptionalAttr<I64Attr>:$comm_block_offset,
Expand All @@ -1761,7 +1762,8 @@ def TPutAsyncOp : PTO_Op<"comm.tput_async", [
let arguments = (ins
PTODpsType:$dst,
PTODpsType:$src,
AsyncSessionType:$session
AsyncSessionType:$session,
DefaultValuedAttr<PTO_DmaEngineAttr, "::mlir::pto::DmaEngine::SDMA">:$dmaEngine
);

let results = (outs AsyncEventType:$event);
Expand All @@ -1782,7 +1784,8 @@ def TGetAsyncOp : PTO_Op<"comm.tget_async", [
let arguments = (ins
PTODpsType:$dst,
PTODpsType:$src,
AsyncSessionType:$session
AsyncSessionType:$session,
DefaultValuedAttr<PTO_DmaEngineAttr, "::mlir::pto::DmaEngine::SDMA">:$dmaEngine
);

let results = (outs AsyncEventType:$event);
Expand Down Expand Up @@ -1933,6 +1936,7 @@ def TBroadcastOp : PTO_Op<"comm.tbroadcast", [
PTODpsType:$ping,
Optional<PTODpsType>:$pong,
Variadic<PTODpsType>:$group,
DefaultValuedAttr<PTO_CollEngineAttr, "::mlir::pto::CollEngine::AIV">:$collEngine,
I32Attr:$root
Comment on lines +1939 to 1940
Copy link
Copy Markdown

@gemini-code-assist gemini-code-assist Bot May 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Inserting the new collEngine attribute before the existing root attribute in the TableGen arguments list breaks backward compatibility for both C++ and Python APIs. Any existing code calling TBroadcastOp with positional arguments will now have their arguments mismatched (e.g., passing the root rank to collEngine). Placing new optional or default-valued attributes at the end of the arguments list preserves the positional argument order and maintains backward compatibility.

      I32Attr:$root,
      DefaultValuedAttr<PTO_CollEngineAttr, "::mlir::pto::CollEngine::AIV">:$collEngine

);
let results = (outs);
Expand All @@ -1950,6 +1954,7 @@ def CommTGatherOp : PTO_Op<"comm.tgather", [
PTODpsType:$ping,
Optional<PTODpsType>:$pong,
Variadic<PTODpsType>:$group,
DefaultValuedAttr<PTO_CollEngineAttr, "::mlir::pto::CollEngine::AIV">:$collEngine,
I32Attr:$root
Comment on lines +1957 to 1958
Copy link
Copy Markdown

@gemini-code-assist gemini-code-assist Bot May 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Inserting the new collEngine attribute before the existing root attribute in the TableGen arguments list breaks backward compatibility for both C++ and Python APIs. Placing new optional or default-valued attributes at the end of the arguments list preserves the positional argument order and maintains backward compatibility.

      I32Attr:$root,
      DefaultValuedAttr<PTO_CollEngineAttr, "::mlir::pto::CollEngine::AIV">:$collEngine

);
let results = (outs);
Expand All @@ -1967,6 +1972,7 @@ def CommTScatterOp : PTO_Op<"comm.tscatter", [
PTODpsType:$ping,
Optional<PTODpsType>:$pong,
Variadic<PTODpsType>:$group,
DefaultValuedAttr<PTO_CollEngineAttr, "::mlir::pto::CollEngine::AIV">:$collEngine,
I32Attr:$root
Comment on lines +1975 to 1976
Copy link
Copy Markdown

@gemini-code-assist gemini-code-assist Bot May 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Inserting the new collEngine attribute before the existing root attribute in the TableGen arguments list breaks backward compatibility for both C++ and Python APIs. Placing new optional or default-valued attributes at the end of the arguments list preserves the positional argument order and maintains backward compatibility.

      I32Attr:$root,
      DefaultValuedAttr<PTO_CollEngineAttr, "::mlir::pto::CollEngine::AIV">:$collEngine

);
let results = (outs);
Expand All @@ -1985,6 +1991,7 @@ def TReduceOp : PTO_Op<"comm.treduce", [
PTODpsType:$recvPing,
Optional<PTODpsType>:$recvPong,
Variadic<PTODpsType>:$group,
DefaultValuedAttr<PTO_CollEngineAttr, "::mlir::pto::CollEngine::AIV">:$collEngine,
PTO_ReduceOpAttr:$reduceOp,
I32Attr:$root
Comment on lines +1994 to 1996
Copy link
Copy Markdown

@gemini-code-assist gemini-code-assist Bot May 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Inserting the new collEngine attribute before the existing reduceOp and root attributes in the TableGen arguments list breaks backward compatibility for both C++ and Python APIs. Placing new optional or default-valued attributes at the end of the arguments list preserves the positional argument order and maintains backward compatibility.

      PTO_ReduceOpAttr:$reduceOp,
      I32Attr:$root,
      DefaultValuedAttr<PTO_CollEngineAttr, "::mlir::pto::CollEngine::AIV">:$collEngine

);
Expand Down
6 changes: 6 additions & 0 deletions include/pto-c/Dialect/PTO.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,12 @@ MLIR_CAPI_EXPORTED int32_t mlirPTOWaitCmpAttrGetValue(MlirAttribute attr);
MLIR_CAPI_EXPORTED bool mlirPTOAttrIsAReduceOpAttr(MlirAttribute attr);
MLIR_CAPI_EXPORTED MlirAttribute mlirPTOReduceOpAttrGet(MlirContext ctx, int32_t value);
MLIR_CAPI_EXPORTED int32_t mlirPTOReduceOpAttrGetValue(MlirAttribute attr);
MLIR_CAPI_EXPORTED bool mlirPTOAttrIsADmaEngineAttr(MlirAttribute attr);
MLIR_CAPI_EXPORTED MlirAttribute mlirPTODmaEngineAttrGet(MlirContext ctx, int32_t value);
MLIR_CAPI_EXPORTED int32_t mlirPTODmaEngineAttrGetValue(MlirAttribute attr);
MLIR_CAPI_EXPORTED bool mlirPTOAttrIsACollEngineAttr(MlirAttribute attr);
MLIR_CAPI_EXPORTED MlirAttribute mlirPTOCollEngineAttrGet(MlirContext ctx, int32_t value);
MLIR_CAPI_EXPORTED int32_t mlirPTOCollEngineAttrGetValue(MlirAttribute attr);
MLIR_CAPI_EXPORTED MlirAttribute mlirPTORoundModeAttrGet(MlirContext ctx, int32_t value);
MLIR_CAPI_EXPORTED bool mlirPTOAttrIsARoundModeAttr(MlirAttribute attr);
MLIR_CAPI_EXPORTED int32_t mlirPTORoundModeAttrGetValue(MlirAttribute attr);
Expand Down
36 changes: 36 additions & 0 deletions lib/Bindings/Python/PTOModule.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,16 @@ static void bindPTOModule(pybind11::module &m) {
.value("Min", mlir::pto::ReduceOp::Min)
.export_values();

py::enum_<mlir::pto::DmaEngine>(m, "DmaEngine")
.value("SDMA", mlir::pto::DmaEngine::SDMA)
.value("URMA", mlir::pto::DmaEngine::URMA)
.export_values();

py::enum_<mlir::pto::CollEngine>(m, "CollEngine")
.value("AIV", mlir::pto::CollEngine::AIV)
.value("CCU", mlir::pto::CollEngine::CCU)
.export_values();

py::enum_<mlir::pto::SyncOpType>(m, "SyncOpType")
.value("TLOAD", mlir::pto::SyncOpType::TLOAD)
.value("TSTORE_ACC", mlir::pto::SyncOpType::TSTORE_ACC)
Expand Down Expand Up @@ -363,6 +373,32 @@ static void bindPTOModule(pybind11::module &m) {
return cls(a);
},
py::arg("cls"), py::arg("value"), py::arg("context") = py::none());

mlir_attribute_subclass(m, "DmaEngineAttr",
[](MlirAttribute a) -> bool {
return mlirPTOAttrIsADmaEngineAttr(a);
})
.def_classmethod(
"get",
[](py::object cls, mlir::pto::DmaEngine value, MlirContext ctx) -> py::object {
MlirAttribute a = mlirPTODmaEngineAttrGet(ctx, static_cast<int32_t>(value));
Comment on lines +383 to +384
Copy link
Copy Markdown

@gemini-code-assist gemini-code-assist Bot May 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Using MlirContext directly in the Python binding lambda prevents the context from being optional (defaulting to the active context in the thread-local stack). If the context is omitted or passed as None in Python, this will either raise a TypeError or pass a null context to mlirPTODmaEngineAttrGet, leading to a segmentation fault when unwrap(ctx) is called. Use DefaultingPyMlirContext instead, which automatically resolves to the default context.

Suggested change
[](py::object cls, mlir::pto::DmaEngine value, MlirContext ctx) -> py::object {
MlirAttribute a = mlirPTODmaEngineAttrGet(ctx, static_cast<int32_t>(value));
[](py::object cls, mlir::pto::DmaEngine value, DefaultingPyMlirContext ctx) -> py::object {
MlirAttribute a = mlirPTODmaEngineAttrGet(ctx.get(), static_cast<int32_t>(value));

if (mlirAttributeIsNull(a)) return py::none();
return cls(a);
},
py::arg("cls"), py::arg("value"), py::arg("context") = py::none());

mlir_attribute_subclass(m, "CollEngineAttr",
[](MlirAttribute a) -> bool {
return mlirPTOAttrIsACollEngineAttr(a);
})
.def_classmethod(
"get",
[](py::object cls, mlir::pto::CollEngine value, MlirContext ctx) -> py::object {
MlirAttribute a = mlirPTOCollEngineAttrGet(ctx, static_cast<int32_t>(value));
Comment on lines +396 to +397
Copy link
Copy Markdown

@gemini-code-assist gemini-code-assist Bot May 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Using MlirContext directly in the Python binding lambda prevents the context from being optional. If the context is omitted or passed as None in Python, this will either raise a TypeError or pass a null context to mlirPTOCollEngineAttrGet, leading to a segmentation fault when unwrap(ctx) is called. Use DefaultingPyMlirContext instead, which automatically resolves to the default context.

Suggested change
[](py::object cls, mlir::pto::CollEngine value, MlirContext ctx) -> py::object {
MlirAttribute a = mlirPTOCollEngineAttrGet(ctx, static_cast<int32_t>(value));
[](py::object cls, mlir::pto::CollEngine value, DefaultingPyMlirContext ctx) -> py::object {
MlirAttribute a = mlirPTOCollEngineAttrGet(ctx.get(), static_cast<int32_t>(value));

if (mlirAttributeIsNull(a)) return py::none();
return cls(a);
},
py::arg("cls"), py::arg("value"), py::arg("context") = py::none());
// [保留 HEAD]: AddressSpaceAttr 定义
mlir_attribute_subclass(
m, "AddressSpaceAttr",
Expand Down
30 changes: 30 additions & 0 deletions lib/CAPI/Dialect/PTO.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -708,6 +708,36 @@ int32_t mlirPTOReduceOpAttrGetValue(MlirAttribute attr) {
return static_cast<int32_t>(a.getValue());
}

bool mlirPTOAttrIsADmaEngineAttr(MlirAttribute attr) {
return mlir::isa<mlir::pto::DmaEngineAttr>(unwrap(attr));
}

MlirAttribute mlirPTODmaEngineAttrGet(MlirContext ctx, int32_t value) {
auto *c = unwrap(ctx);
return wrap(mlir::pto::DmaEngineAttr::get(
c, static_cast<mlir::pto::DmaEngine>(value)));
}

int32_t mlirPTODmaEngineAttrGetValue(MlirAttribute attr) {
auto a = mlir::cast<mlir::pto::DmaEngineAttr>(unwrap(attr));
return static_cast<int32_t>(a.getValue());
}

bool mlirPTOAttrIsACollEngineAttr(MlirAttribute attr) {
return mlir::isa<mlir::pto::CollEngineAttr>(unwrap(attr));
}

MlirAttribute mlirPTOCollEngineAttrGet(MlirContext ctx, int32_t value) {
auto *c = unwrap(ctx);
return wrap(mlir::pto::CollEngineAttr::get(
c, static_cast<mlir::pto::CollEngine>(value)));
}

int32_t mlirPTOCollEngineAttrGetValue(MlirAttribute attr) {
auto a = mlir::cast<mlir::pto::CollEngineAttr>(unwrap(attr));
return static_cast<int32_t>(a.getValue());
}

MlirAttribute mlirPTOTileBufConfigAttrGet(MlirContext ctx,
MlirAttribute bLayout,
MlirAttribute sLayout,
Expand Down
Loading
Loading