8000 [MLIR][OpenMP] Add `omp.private` op · llvm/llvm-project@0d165e0 · GitHub
[go: up one dir, main page]
Skip to content

    • Commit 0d165e0

    Browse files
    ergawyergawy
    committed
    [MLIR][OpenMP] Add omp.private op
    This PR adds a new op to the OpenMP dialect: `PrivateClauseOp`. This op will be later used to model `[first]private` clauses for differnt OpenMP directives. This is part of productizing the "delayed privatization" PoC wich can be found in #79862.
    1 parent 0c63453 commit 0d165e0

    File tree

    4 files changed

    +243
    -1
    lines changed

    4 files changed

    +243
    -1
    lines changed

    ‎mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

    Lines changed: 92 additions & 1 deletion
    Original file line numberDiff line numberDiff line change
    @@ -133,6 +133,97 @@ def DeclareTargetAttr : OpenMP_Attr<"DeclareTarget", "declaretarget"> {
    133133
    let assemblyFormat = "`<` struct(params) `>`";
    134134
    }
    135135

    136+
    //===----------------------------------------------------------------------===//
    137+
    // 2.19.4 Data-Sharing Attribute Clauses
    138+
    //===----------------------------------------------------------------------===//
    139+
    140+
    def DataSharingTypePrivate : I32EnumAttrCase<"Private", 0, "private">;
    141+
    def DataSharingTypeFirstPrivate : I32EnumAttrCase<"FirstPrivate", 1, "firstprivate">;
    142+
    143+
    def DataSharingClauseType : I32EnumAttr<
    144+
    "DataSharingClauseType",
    145+
    "Type of a data-sharing clause",
    146+
    [DataSharingTypePrivate, DataSharingTypeFirstPrivate]> {
    147+
    let genSpecializedAttr = 0;
    148+
    let cppNamespace = "::mlir::omp";
    149+
    }
    150+
    151+
    def DataSharingClauseTypeAttr : EnumAttr<
    152+
    OpenMP_Dialect, DataSharingClauseType, "data_sharing_type"> {
    153+
    let assemblyFormat = "`{` `type` `=` $value `}`";
    154+
    }
    155+
    156+
    def PrivateClauseOp : OpenMP_Op<"private", [IsolatedFromAbove]> {
    157+
    let summary = "Provides declaration of [first]private logic.";
    158+
    let description = [{
    159+
    This operation provides a declaration of how to implement the
    160+
    [first]privatization of a variable. The dialect users should provide
    161+
    information about how to create an instance of the type in the alloc region
    162+
    and how to initialize the copy from the original item in the copy region.
    163+
    164+
    Examples:
    165+
    ---------
    166+
    * `private(x)` would be emitted as:
    167+
    ```mlir
    168+
    omp.private {type = private} @x.privatizer : !fir.ref<i32> alloc {
    169+
    ^bb0(%arg0: !fir.ref<i32>):
    170+
    %0 = ... allocate proper memory for the private clone ...
    171+
    omp.yield(%0 : !fir.ref<i32>)
    172+
    }
    173+
    ```
    174+
    175+
    * `firstprivate(x)` would be emitted as:
    176+
    ```mlir
    177+
    omp.private {type = firstprivate} @x.privatizer : !fir.ref<i32> alloc {
    178+
    ^bb0(%arg0: !fir.ref<i32>):
    179+
    %0 = ... allocate proper memory for the private clone ...
    180+
    omp.yield(%0 : !fir.ref<i32>)
    181+
    } copy {
    182+
    ^bb0(%arg0: !fir.ref<i32>, %arg1: !fir.ref<i32>):
    183+
    // %arg0 is the original host variable. Same as for `alloc`.
    184+
    // %arg1 represents the memory allocated in `alloc`.
    185+
    ... copy from host to the privatized clone ....
    186+
    omp.yield(%arg1 : !fir.ref<i32>)
    187+
    }
    188+
    ```
    189+
    190+
    There are no restrictions on the body except for:
    191+
    - The `alloc` region has a single argument.
    192+
    - The `copy` region has 2 arguments.
    193+
    - Both regions are terminated by `omp.yield` ops.
    194+
    The above restrictions and other obvious restrictions (e.g. verifying the
    195+
    type of yielded values) are verified by the custom op verifier. The actual
    196+
    contents of the blocks inside both regions are not verified.
    197+
    198+
    Instances of this op would then be used by ops that model directives that
    199+
    accept data-sharing attribute clauses.
    200+
    201+
    The $sym_name attribute provides a symbol by which the privatizer op can be
    202+
    referenced by other dialect ops.
    203+
    204+
    The $type attribute is the type of the value being privatized.
    205+
    206+
    The $data_sharing_type attribute specifies whether privatizer corresponds
    207+
    to a `private` or a `firstprivate` clause.
    208+
    }];
    209+
    210+
    let arguments = (ins SymbolNameAttr:$sym_name,
    211+
    TypeAttrOf<AnyType>:$type,
    212+
    DataSharingClauseTypeAttr:$data_sharing_type);
    213+
    214+
    let regions = (region MinSizedRegion<1>:$alloc_region,
    215+
    AnyRegion:$copy_region);
    216+
    217+
    let assemblyFormat = [{
    218+
    $data_sharing_type $sym_name `:` $type
    219+
    `alloc` $alloc_region
    220+
    (`copy` $copy_region^)?
    221+
    attr-dict
    222+
    }];
    223+
    224+
    let hasVerifier = 1;
    225+
    }
    226+
    136227
    //===----------------------------------------------------------------------===//
    137228
    // 2.6 parallel Construct
    138229
    //===----------------------------------------------------------------------===//
    @@ -612,7 +703,7 @@ def SimdLoopOp : OpenMP_Op<"simdloop", [AttrSizedOperandSegments,
    612703
    def YieldOp : OpenMP_Op<"yield",
    613704
    [Pure, ReturnLike, Terminator,
    614705
    ParentOneOf<["WsLoopOp", "ReductionDeclareOp",
    615-
    "AtomicUpdateOp", "SimdLoopOp"]>]> {
    706+
    "AtomicUpdateOp", "SimdLoopOp", "PrivateClauseOp"]>]> {
    616707
    let summary = "loop yield and termination operation";
    617708
    let description = [{
    618709
    "omp.yield" yields SSA values from the OpenMP dialect op region and

    ‎mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

    Lines changed: 67 additions & 0 deletions
    Original file line numberDiff line numberDiff line change
    @@ -1594,6 +1594,73 @@ LogicalResult DataBoundsOp::verify() {
    15941594
    return success();
    15951595
    }
    15961596

    1597+
    LogicalResult PrivateClauseOp::verify() {
    1598+
    Type symType = getType();
    1599+
    1600+
    auto verifyTerminator = [&](Operation *terminator) -> LogicalResult {
    1601+
    if (!terminator->hasSuccessors() && !llvm::isa<YieldOp>(terminator))
    1602+
    return mlir::emitError(terminator->getLoc())
    1603+
    << "expected exit block terminator to be an `omp.yield` op.";
    1604+
    1605+
    YieldOp yieldOp = llvm::cast<YieldOp>(terminator);
    1606+
    TypeRange yieldedTypes = yieldOp.getResults().getTypes();
    1607+
    1608+
    if (yieldedTypes.size() == 1 && yieldedTypes.front() == symType)
    1609+
    return success();
    1610+
    1611+
    auto error = mlir::emitError(yieldOp.getLoc())
    1612+
    << "Invalid yielded value. Expected type: " << symType
    1613+
    << ", got: ";
    1614+
    1615+
    if (yieldedTypes.empty())
    1616+
    error << "None";
    1617+
    else
    1618+
    error << yieldedTypes;
    1619+
    1620+
    return error;
    1621+
    };
    1622+
    1623+
    auto verifyRegion = [&](Region &region, unsigned expectedNumArgs,
    1624+
    StringRef regionName) -> LogicalResult {
    1625+
    assert(!region.empty());
    1626+
    1627+
    if (region.getNumArguments() != expectedNumArgs)
    1628+
    return mlir::emitError(region.getLoc())
    1629+
    << "`" << regionName << "`: "
    1630+
    << "expected " << expectedNumArgs
    1631+
    << " region arguments, got: " << region.getNumArguments();
    1632+
    1633+
    for (Block &block : region) {
    1634+
    // MLIR will verify the absence of the terminator for us.
    1635+
    if (!block.mightHaveTerminator())
    1636+
    continue;
    1637+
    1638+
    if (failed(verifyTerminator(block.getTerminator())))
    1639+
    return failure();
    1640+
    }
    1641+
    1642+
    return success();
    1643+
    };
    1644+
    1645+
    if (failed(verifyRegion(getAllocRegion(), /*expectedNumArgs=*/1, "alloc")))
    1646+
    return failure();
    1647+
    1648+
    DataSharingClauseType dsType = getDataSharingType();
    1649+
    1650+
    if (dsType == DataSharingClauseType::Private && !getCopyRegion().empty())
    1651+
    return emitError("`private` clauses require only an `alloc` region.");
    1652+
    1653+
    if (dsType == DataSharingClauseType::FirstPrivate && getCopyRegion().empty())
    1654+
    return emitError(
    1655+
    "`firstprivate` clauses require both `alloc` and `copy` regions.");
    1656+
    1657+
    if (dsType == DataSharingClauseType::FirstPrivate &&
    1658+
    failed(verifyRegion(getCopyRegion(), /*expectedNumArgs=*/2, "copy")))
    1659+
    return failure();
    1660+
    1661+
    return success();
    1662+
    }
    1663+
    15971664
    #define GET_ATTRDEF_CLASSES
    15981665
    #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
    15991666

    ‎mlir/test/Dialect/OpenMP/invalid.mlir

    Lines changed: 63 additions & 0 deletions
    Original file line numberDiff line numberDiff line change
    @@ -1738,3 +1738,66 @@ func.func @omp_distribute(%data_var : memref<i32>) -> () {
    17381738
    "omp.terminator"() : () -> ()
    17391739
    }) : (memref<i32>) -> ()
    17401740
    }
    1741+
    1742+
    // -----
    1743+
    1744+
    omp.private {type = private} @x.privatizer : i32 alloc {
    1745+
    ^bb0(%arg0: i32):
    1746+
    %0 = arith.constant 0.0 : f32
    1747+
    // expected-error @below {{Invalid yielded value. Expected type: 'i32', got: 'f32'}}
    1748+
    omp.yield(%0 : f32)
    1749+
    }
    1750+
    1751+
    // -----
    1752+
    1753+
    omp.private {type = private} @x.privatizer : i32 alloc {
    1754+
    ^bb0(%arg0: i32):
    1755+
    // expected-error @below {{Invalid yielded value. Expected type: 'i32', got: None}}
    1756+
    omp.yield
    1757+
    }
    1758+
    1759+
    // -----
    1760+
    1761+
    omp.private {type = private} @x.privatizer : i32 alloc {
    1762+
    ^bb0(%arg0: i32):
    1763+
    // expected-error @below {{expected exit block terminator to be an `omp.yield` op.}}
    1764+
    omp.terminator
    1765+
    }
    1766+
    1767+
    // -----
    1768+
    1769+
    // expected-error @below {{`alloc`: expected 1 region arguments, got: 2}}
    1770+
    omp.private {type = private} @x.privatizer : f32 alloc {
    1771+
    ^bb0(%arg0: f32, %arg1: f32):
    1772+
    omp.yield(%arg0 : f32)
    1773+
    }
    1774+
    1775+
    // -----
    1776+
    1777+
    // expected-error @below {{`copy`: expected 2 region arguments, got: 1}}
    1778+
    omp.private {type = firstprivate} @x.privatizer : f32 alloc {
    1779+
    ^bb0(%arg0: f32):
    1780+
    omp.yield(%arg0 : f32)
    1781+
    } copy {
    1782+
    ^bb0(%arg0: f32):
    1783+
    omp.yield(%arg0 : f32)
    1784+
    }
    1785+
    1786+
    // -----
    1787+
    1788+
    // expected-error @below {{`private` clauses require only an `alloc` region.}}
    1789+
    omp.private {type = private} @x.privatizer : f32 alloc {
    1790+
    ^bb0(%arg0: f32):
    1791+
    omp.yield(%arg0 : f32)
    1792+
    } copy {
    1793+
    ^bb0(%arg0: f32, %arg1 : f32):
    1794+
    omp.yield(%arg0 : f32)
    1795+
    }
    1796+
    1797+
    // -----
    1798+
    1799+
    // expected-error @below {{`firstprivate` clauses require both `alloc` and `copy` regions.}}
    1800+
    omp.private {type = firstprivate} @x.privatizer : f32 alloc {
    1801+
    ^bb0(%arg0: f32):
    1802+
    omp.yield(%arg0 : f32)
    1803+
    }

    ‎mlir/test/Dialect/OpenMP/roundtrip.mlir

    Lines changed: 21 additions & 0 deletions
    Original file line numberDiff line numberDiff line change
    @@ -0,0 +1,21 @@
    1+
    // RUN: mlir-opt -verify-diagnostics %s | mlir-opt | FileCheck %s
    2+
    3+
    // CHECK: omp.private {type = private} @x.privatizer : !llvm.ptr alloc {
    4+
    omp.private {type = private} @x.privatizer : !llvm.ptr alloc {
    5+
    // CHECK: ^bb0(%arg0: {{.*}}):
    6+
    ^bb0(%arg0: !llvm.ptr):
    7+
    omp.yield(%arg0 : !llvm.ptr)
    8+
    }
    9+
    10+
    // CHECK: omp.private {type = firstprivate} @y.privatizer : !llvm.ptr alloc {
    11+
    omp.private {type = firstprivate} @y.privatizer : !llvm.ptr alloc {
    12+
    // CHECK: ^bb0(%arg0: {{.*}}):
    13+
    ^bb0(%arg0: !llvm.ptr):
    14+
    omp.yield(%arg0 : !llvm.ptr)
    15+
    // CHECK: } copy {
    16+
    } copy {
    17+
    // CHECK: ^bb0(%arg0: {{.*}}, %arg1: {{.*}}):
    18+
    ^bb0(%arg0: !llvm.ptr, %arg1: !llvm.ptr):
    19+
    omp.yield(%arg0 : !llvm.ptr)
    20+
    }
    21+

    0 commit comments

    Comments
     (0)
    0