Skip to content

Formalize syntax and evaluation of StableHLO programs #484

Open
Open
Formalize syntax and evaluation of StableHLO programs#484
Labels
Spec

Description

@burmako
burmako
opened on Nov 8, 2022

We've made significant progress with the spec in the last two months, and it looks like soon we will have all of the core ops specced after the existing HLO semantics. There's a bunch of ops in the PR queue right now, and about a dozen more distribution and elementwise ops to go, but finishing all that looks like a matter of several weeks at this point thanks to the great work done by the team:

Category Ops
Control Flow case, if, while
Data Movement broadcast_in_dim, concatenate, gather, pad, reshape, reverse, scatter, slice, sort, transpose
Distribution after_all, all_gather, all_reduce, all_to_all, collective_permute, infeed, outfeed, recv, reduce_scatter, replica_id, send
Elementwise abs, add, and, atan2, bitcast_convert, cbrt, ceil, clamp, compare, complex, convert, cosine, count_leading_zeros, divide, exponential, exponential_minus_one, floor, imag, is_finite, log, log_plus_one, logistic, map, maximum, minimum, multiply, negate, not, or, popcnt, power, real, reduce_precision, remainder, round_nearest_afz, round_nearest_even, rsqrt, select, shift_left, shift_right_arithmetic, shift_right_logical, sign, sine, sqrt, subtract, tanh, xor
Extensibility custom_call, get_tuple_element, tuple
Miscellaneous batch_norm_grad, batch_norm_inference, batch_norm_training, cholesky, constant, fft, iota, optimization_barrier, return, rng, rng_bit_generator, triangular_solve
Reduction convolution, dot_general, reduce, reduce_window, select_and_scatter

This will make it possible to do something cool that wasn't possible before - express RFCs as diffs to the spec - meaning that the work on evolution beyond what's currently in HLO would become more precise:

Category Ops
Dynamism compute_reshape_shape, cstr_reshapable, dynamic_broadcast_in_dim, dynamic_conv, dynamic_gather, dynamic_iota, dynamic_pad, dynamic_reshape, dynamic_slice, dynamic_update_slice, get_dimension_size, real_dynamic_slice, set_dimension_size
Quantization uniform_dequantize, uniform_quantize + <new element types>
Sparsity <new tensor types>

To enable this, we're missing one key piece - framing the spec to provide self-contained syntax for StableHLO programs and self-contained rules for their evaluation. So far, we've been half-handwaving this / half-piggybacking on MLIR, but I think we're fairly close to have something truly self-contained, which will allow us to do really cool stuff.

Activity

  1. burmako
    added
    Spec
    on Nov 8, 2022
  2. burmako
    self-assigned this
    on Nov 8, 2022

  3. burmako commented on Nov 30, 2022

    @burmako
    burmako
    on Nov 30, 2022
    Author

    Feedback from #634:

    1. What happens if some of the ops in a function aren't reachable from its return op?
    2. We could use topological order as a vehicle for formalizing execution order.
    3. Perhaps provide more details on the purpose of replicas and partitions?
    4. What happens to replica_id if the program doesn't run in parallel.
  4. burmako
    added a commit that references this issue on Nov 30, 2022

    Add spec for parallel execution semantics (#634)

    336eac4
  5. burmako
    added a commit that references this issue on Dec 27, 2022

    First step towards turning spec.md into a programming language spec (#…

    57c6015
  6. burmako
    added this to Documentationon Feb 6, 2023
  7. burmako
    moved this to Todo in Documentationon Feb 19, 2023

  8. sdasgup3 commented on Mar 10, 2023

    @sdasgup3
    sdasgup3
    on Mar 10, 2023
    Member

    In #995 (comment), we informally talked about the scoping rules. Here is an attempt to formalize the same.

    An SSA value in StableHLO is either an argument of a block or result of an operation. Each SSA value has zero or more uses. An operation defines zero or more SSA value as results, and has zero or more operands: these are uses of SSA value, which can be the results of other operations or block arguments.

    For an arbitrary region R, value V, and block B,

    • parent(R): Parent operation this region is attached to.
    • parent(B): Region to which the block B belongs.
    • parent_region(op): Region to which the operation op belongs. For op = func.func, returns nullptr.
    • parent_region(R) = parent_region(parent(R))
    • is_defined(V, R) = (SSA value V is either defined by an operation in block B or a block argument of B) and parent(B) = R.

    For a use U of SSA value V, we define a function,in_scope(U, R) -> bool, which evaluates to true if U can be used as operand of an operation op such that parent_region(op) = R.

    in_scope(U, R) is given by

    • false, if R = nullptr,
    • is_defined(V, R) or in_scope(U, parent_region(R)), otherwise.

    cc @burmako

  9. burmako commented on May 24, 2023

    @burmako
    burmako
    on May 24, 2023
    Author

    There are some ops which express result[result_index] through operand[operand_index], and some ops which say stuff like result[i0, ..., ik,..., iR-1] = operand[i0, ..., ik',..., iR-1] where ik + ik' = dk - 1 for all dimensions k in dimensions.

    Might be interesting to standardize on one way of doing things. For example, the latter example (for ReverseOp) could be reformulated as:

    `result[result_index] = operand[operand_index]` where:
  * `operand_index[d] = dim(operand, d) - result_index[d]` if `d` in `dimensions`.
  * `operand_index[d] = result_index[d]` otherwise.
  10. burmako
    added a commit that references this issue on Jun 2, 2023

    Align the notation of spec and interpreter for Data Movement operatio…

    1677e2e
  11. burmako
    removed their assignment
    on Jun 6, 2023
  12. burmako
    moved this from Todo to No status in Documentationon Jun 6, 2023

  13. burmako commented on Jun 6, 2023

    @burmako
    burmako
    on Jun 6, 2023
    Author

    After #1556 landed, I'm reasonably happy with how the current version of the spec formalizes StableHLO syntax. There's more work to be done, but at least the notation is roughly consistent now across ops and is explained in the specification. Unassigning this ticket for now, to focus on the dynamism RFC.

  14. joelberkeley commented on Apr 5, 2025

    @joelberkeley
    joelberkeley
    on Apr 5, 2025

    What's the situation regarding send/recv ops? Afaik the semantics wrt channel IDs - specifically how to create channels and get their ID - is not well-defined yet. Is that correct?

  15. sdasgup3 commented on Apr 7, 2025

    @sdasgup3
    sdasgup3
    on Apr 7, 2025
    Member

    That's correct. The StableHLO specification currently defines when a channel_id is required by certain operations but not the mechanism for creating channels or obtaining their IDs.

    The focus within StableHLO is on the operation's semantics:

    For collective operations within a process group, a channel (identified by a positive `si64` channel_id) is required specifically for cross-partition communication (when processes involved are on different partitions). Communication that is purely cross-replica (within the same partition) does not necessitate a channel ID according to the spec.

    How frameworks or runtime systems actually establish these communication channels and generate/manage the corresponding channel_id values is considered an implementation detail while drafting the specification.

  16. joelberkeley commented on Apr 7, 2025

    @joelberkeley
    joelberkeley
    on Apr 7, 2025

    Also appears to be beyond the scope of PJRT. I'm very unclear on how to create a channel. Should I use NCCL directly? I'm unclear what CUDA's PJRT plugin does already behind the scenes, and the semantics of using NCCL directly, alongside PJRT/stablehlo.

    p.s. can move to discord if this isn't the right place for this

  17. sdasgup3 commented on Apr 7, 2025

    @sdasgup3
    sdasgup3
    on Apr 7, 2025
    Member

    can move to discord if this isn't the right place for this

    I agree that would be the right choice (for more eyes with CUDA PJRT expertise)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Spec

    Type

    No type

    Projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions