Ngôn ngữ "sdy"

Phương ngữ Shardy (SDY) xác định một đại diện phân đoạn tensor dựa trên trục và các thành phần API bổ sung để đính kèm các phân đoạn vào tensor.

Hoạt động tính toán

sdy.constant (sdy::ConstantOp)

Toán tử hằng

Tạo tensor output từ một hằng số value.

Xem: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

Ví dụ:

%output = sdy.constant dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>

Đặc điểm: AlwaysSpeculatableImplTrait

Giao diện: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Hiệu ứng: MemoryEffects::Effect{}

Thuộc tính:

Kết quả:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Toán tử cạnh luồng dữ liệu.

Cú pháp:

operation ::= `sdy.data_flow_edge` $input (`sharding````=``` $sharding^)? attr-dict `:` type($result)

Cạnh luồng dữ liệu của một số op X xác định cầu nối giữa một tập hợp các nguồn (mỗi nguồn là một toán hạng của X hoặc một toán hạng của phần tử kết thúc khối của X) và một tập hợp các mục tiêu (mỗi mục tiêu là kết quả của X hoặc một đối số khối của X), sao cho tất cả các nguồn và mục tiêu phải được phân đoạn theo cùng một cách.

Một hoạt động có thể có nhiều cạnh luồng dữ liệu trực giao với nhau.

Ví dụ:

  y_0, ..., y_n = while (x_0, ..., x_n)
                  ((pred_arg_0,... , pred_arg_n) { ... })
                  ((body_arg_0,..., body_arg_n) {
                    ...
                    return return_value_0, ..., return_value_n
                  })

Toán tử while này có n cạnh luồng dữ liệu, cạnh luồng dữ liệu thứ i nằm giữa nguồn x_i, return_value_i và đích y_i, pred_arg_i, body_arg_i.

sdy.data_flow_edge lấy làm đầu vào mục tiêu gốc của một cạnh (có thể là bất kỳ mục tiêu nào, nhưng tốt nhất là kết quả toán tử thay vì đối số khối), không được có bất kỳ mục đích sử dụng nào khác. Toán tử này không thuần tuý vì có thể lấy một giá trị đầu vào ban đầu không có bất kỳ giá trị sử dụng nào.

sdy.data_flow_edge cũng chứa một phân đoạn không bắt buộc cho tất cả các mục tiêu của cạnh và phân đoạn đó phải được cập nhật thay vì phân đoạn của mục tiêu (nếu có thể đính kèm) trong quá trình truyền tải. Điều này rất hữu ích khi một toán tử có nhiều cạnh, vì việc này sẽ hiệu quả hơn nhiều khi:

• truyền qua từng cạnh riêng biệt.
• cập nhật việc phân đoạn của từng cạnh riêng biệt thay vì tất cả các mục tiêu cùng một lúc (ví dụ: một hoạt động có một TensorShardingPerValueAttr không thể thay đổi để phân đoạn kết quả).
• thêm từng cạnh vào danh sách công việc riêng biệt khi phân đoạn của một nguồn đã thay đổi.

Quá trình truyền sẽ truyền các phân đoạn giữa tất cả các nguồn và mục tiêu của sdy.data_flow_edge như thể đó là một hoạt động thông thường với các nguồn là toán hạng và nhắm mục tiêu là kết quả và một danh tính sdy.op_sharding_rule. Điều đó có nghĩa là quá trình truyền về phía trước là từ nguồn đến đích và quá trình truyền về phía sau là từ đích đến nguồn.

Chúng tôi không cho phép thao tác đầu vào của sdy.data_flow_edge được xác định bằng thao tác SdyDialect, vì vậy, chúng ta có thể giả định rằng thao tác đầu vào này được xác định bằng thao tác có thuộc tính sdy.sharding chưa đăng ký.

Đặc điểm: SameOperandsAndResultType

Giao diện: InferTypeOpInterface

Thuộc tính:

Toán hạng:

Kết quả:

sdy.manual_computation (sdy::ManualComputationOp)

Hoạt động chạy song song nhiều thiết bị bằng tập hợp thủ công

Cú pháp:

operation ::= `sdy.manual_computation` `(`operands`)`
              `in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)
              `out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)
              `manual_axes````=```$manual_axes
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:`
              functional-type(operands, results)

Chuyển đến một vùng được viết bằng mã cục bộ trên mỗi thiết bị với các tập hợp rõ ràng, trong đó các hình dạng logic khớp với các hình dạng vùng đệm vật lý cục bộ trên mỗi thiết bị và các tập hợp tương ứng chính xác với giao tiếp vật lý trên nhiều thiết bị.

Thân xe là cục bộ so với manual_axes. Quá trình truyền sẽ diễn ra thông qua phần nội dung trên bất kỳ trục tự do nào – những trục không có trong danh sách manual_ax.

Trait: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Thuộc tính:

Toán hạng:

Kết quả:

sdy.mesh (sdy::MeshOp)

Lưới được đặt tên

Cú pháp:

operation ::= `sdy.mesh` $sym_name `=` $mesh attr-dict

Xác định một lưới mới được đặt tên. Tất cả các lưới trong một mô-đun phải có cùng số lượng thiết bị (ngoại trừ các lưới có một device_id). Lưới là một phép toán Symbol xuất hiện trong SymbolTable của mô-đun và có thể được tham chiếu bằng name của mô-đun đó.

Đặc điểm: HasParent<ModuleOp>

Giao diện: Symbol

Thuộc tính:

sdy.named_computation (sdy::NamedComputationOp)

Toán tử tính toán được đặt tên

Cú pháp:

operation ::= `sdy.named_computation` `<`$name`>` `` `(` $operands `)`
              (`in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)^)?
              (`out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)^)?
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:` functional-type($operands, results)

Nhóm một phép tính (ví dụ: một khối phép toán) và đặt tên cho phép tính đó. Quá trình truyền sẽ diễn ra vào/ra khỏi vùng này như thể mọi thứ đều được cùng dòng.

Bạn có thể dùng phương thức này để xử lý việc truyền thông qua các lệnh gọi đến các hàm khác. Mọi người dùng Shardy đều phải viết một thẻ nhập/xuất chuyển đổi các thao tác gọi thành thao tác sdy.named_computation, sao chép/sao chép nội dung của hàm được gọi vào nội dung của named_computation.

Loại của mỗi đối số khối và giá trị trả về trong vùng phải giống với loại toán hạng và loại kết quả của toán tử.

Ví dụ:

%1 = sdy.named_computation<"foo">(%0) (%arg1: tensor<16x32xf32>) {
  sdy.return %arg1 : tensor<16x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

Trait: IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Giao diện: ConditionallySpeculatable, ShardableDataFlowOpInterface

Thuộc tính:

Toán hạng:

Kết quả:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Thao tác rào cản truyền tải

Cú pháp:

operation ::= `sdy.propagation_barrier` $input `allowed_direction````=```$allowed_direction attr-dict `:` type($input)

Toán tử này hoạt động giống như toán tử nhận dạng, xuất ra cùng một giá trị đã lấy làm đầu vào. Nhưng về mặt truyền tải, điều này sẽ chỉ cho phép truyền tải chảy qua theo một hướng nhất định.

Điều này ngăn việc phân đoạn được truyền giữa các lần sử dụng kết quả của toán tử rào cản và toán hạng của toán tử đó.

• FORWARD có nghĩa là các phần phân đoạn chỉ có thể chuyển từ toán hạng đến kết quả.
• BACKWARD có nghĩa là các phần phân đoạn chỉ có thể chuyển từ kết quả đến toán hạng.
• NONE có nghĩa là không thể phân đoạn nào có thể truyền qua cơ chế này.
• Không thể chỉ định BOTH vì toán tử này sẽ thừa.

Trait: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

Giao diện: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Hiệu ứng: MemoryEffects::Effect{}

Thuộc tính:

Toán hạng:

Kết quả:

sdy.reshard (sdy::ReshardOp)

Phân đoạn lại một tensor thành một phân đoạn khác

Cú pháp:

operation ::= `sdy.reshard` $input $sharding attr-dict `:` type($result)

Phân đoạn lại tensor đầu vào bằng phương thức phân đoạn đã chỉ định, khác với phương thức phân đoạn hiện có của tensor đầu vào.

Cả ShardingConstraintOp và ReshardOp đều đính kèm một phân đoạn vào một tensor. Thời gian hoạt động của chúng là:

1. Trước khi truyền phân đoạn, người dùng sẽ thêm ShardingConstraintOp.
2. Quá trình truyền tải phân đoạn sẽ sử dụng ShardingConstraintOp. Không có ShardingConstraintOp trong kết quả của quá trình truyền phân đoạn. Thay vào đó, bạn có thể thêm ReshardOp nếu cần.
3. Trình phân vùng chuyển đổi ReshardOp thành một hoạt động tập thể (hoặc một hoạt động nhận dạng). Không được có ReshardOp trong kết quả của trình phân vùng.

// TODO(b/331680067). Thêm mẫu chuẩn hoá để xoá // hoạt động phân phát lại không cần thiết.

Thuộc tính: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

Giao diện: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Hiệu ứng: MemoryEffects::Effect{}

Thuộc tính:

Toán hạng:

Kết quả:

sdy.return (sdy::ReturnOp)

Thao tác sdy.return chấm dứt các vùng được đính kèm vào các thao tác dựa trên vùng sdy và mọi thao tác dựa trên vùng Shardy khác. Đây là một hàm biến đổi: hàm này lấy danh sách các giá trị có kiểu bất kỳ làm đối số (nhưng cùng loại, ví dụ: AnyTensor) và do đó có thể được sử dụng lại ở nhiều cấp của ngăn xếp IR Shardy.

Cú pháp:

operation ::= `sdy.return` attr-dict ($results^ `:` type($results))?

Đặc điểm: AlwaysSpeculatableImplTrait, Terminator

Giao diện: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Hiệu ứng: MemoryEffects::Effect{}

Toán hạng:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Ràng buộc một tensor với phân đoạn đã chỉ định

Cú pháp:

operation ::= `sdy.sharding_constraint` $input $sharding attr-dict `:` type($result)

Đính kèm một phân đoạn vào một tensor trung gian (ví dụ: kết quả của matmul) để cho biết đây là cách phân đoạn tensor đó hoặc một tập hợp con của các mục đích sử dụng.

Nếu phân đoạn có các phương diện mở và trục không ràng buộc, thì điều đó có nghĩa là tensor có thể được phân đoạn thêm dọc theo các phương diện mở.

Thao tác này có thể:

• Không có mục đích sử dụng (dangling) – tức là việc phân đoạn đính kèm là cách phân đoạn chính tensor đầu vào.
• Có trường hợp sử dụng – nghĩa là phân đoạn đính kèm là cách phân đoạn các trường hợp sử dụng của toán tử ràng buộc phân đoạn, trong khi các trường hợp sử dụng khác của tensor đầu vào có thể có phân đoạn khác (nếu tensor đầu vào không có trường hợp sử dụng nào khác thì hành vi sẽ giống như trường hợp không có trường hợp sử dụng).

Đặc điểm: Elementwise, SameOperandsAndResultType

Giao diện: InferTypeOpInterface

Thuộc tính:

Toán hạng:

Kết quả:

sdy.sharding_group (sdy::ShardingGroupOp)

Thao tác phân đoạn nhóm

Cú pháp:

operation ::= `sdy.sharding_group` $input `group_id````=```$group_id attr-dict `:` type($input)

Toán tử này cung cấp một giao diện để gán tensor cho các nhóm phân đoạn ( các nhóm tensor sẽ được thực thi để có các phân đoạn giống hệt nhau). Trong quá trình truyền tải, ngay khi một phần tử nhóm được phân đoạn, tất cả các thành phần khác sẽ được phân đoạn theo đúng cách. Thao tác này lấy mã nhận dạng nhóm đối số và không trả về kết quả nào, nhưng thay vào đó sẽ sửa đổi nội dung trình bày nhóm phân đoạn nội bộ để thêm tensor đầu vào vào nhóm có mã nhận dạng đã cho.

Thuộc tính:

Toán hạng:

Thuộc tính

AxisRefAttr

Tham chiếu đến một trục đầy đủ hoặc một trục phụ được chia

Cú pháp:

#sdy.axis_ref<
  ::llvm::StringRef,   # name
  SubAxisInfoAttr   # sub_axis_info
>

Các thông số:

DimMappingAttr

Danh sách chỉ số nhân tố cho một phương diện

Tất cả chỉ mục nhân tố phải nằm trong phạm vi [0, num_factors) và danh sách trống cho biết đây là mối liên kết rỗng (được phân tích cú pháp/in bằng *), tức là phương diện không được liên kết với bất kỳ nhân tố nào.

Các thông số:

DimensionShardingAttr

Phân đoạn phương diện

Danh sách tên trục để phân đoạn một phương diện tensor từ chính đến phụ, một giá trị boolean cho biết liệu phương diện có thể được phân đoạn thêm hay không và một số nguyên tuỳ chọn biểu thị mức độ ưu tiên của việc phân đoạn phương diện này. Mức độ ưu tiên này sẽ được tuân thủ trong quá trình truyền tải phân đoạn. Mức độ ưu tiên bắt nguồn từ các chú thích phân đoạn người dùng và giá trị thấp hơn biểu thị mức độ ưu tiên cao hơn. Mức độ ưu tiên cao nhất được giả định khi thiếu mức độ ưu tiên trong chú thích.

Các thông số:

ManualAxesAttr

Cú pháp:

#sdy.manual_axes<
  ::llvm::ArrayRef<StringAttr>   # value
>

Các thông số:

MeshAttr

Lưới trục và danh sách thiết bị

Cú pháp:

#sdy.mesh<
  ::llvm::ArrayRef<MeshAxisAttr>,   # axes
  ::llvm::ArrayRef<int64_t>   # device_ids
>

Lưới là danh sách các trục và danh sách mã thiết bị (không bắt buộc) chỉ định thứ tự thiết bị.

Nếu danh sách trục trống, thì lưới sẽ có một trục ẩn không tên có kích thước 1. Trong trường hợp này, nếu danh sách mã thiết bị không được cung cấp, thì danh sách mã thiết bị ngầm ẩn sẽ là [0]; nếu một danh sách mã thiết bị được cung cấp, danh sách đó phải chứa một số nguyên duy nhất của bất kỳ giá trị không âm nào. Chúng tôi gọi đây là trường hợp phân đoạn tối đa.

Đối với tất cả các trường hợp không phân đoạn tối đa, nếu bạn chỉ định danh sách mã thiết bị, thì tích của các kích thước trục phải khớp với số lượng thiết bị. Nếu bạn không chỉ định danh sách mã thiết bị, thì danh sách mã thiết bị ngầm ẩn sẽ là iota(product(axes)). Để đơn giản, chúng tôi cũng không cho phép chỉ định danh sách mã thiết bị giống với iota(product(axes)); trong trường hợp này, bạn không nên chỉ định danh sách mã thiết bị.

Dưới đây là một số ví dụ về lưới:

• Lưới trống đại diện cho lưới giữ chỗ có thể được thay thế trong quá trình truyền: <[]>
• Một lưới có trục không tên và mã thiết bị rõ ràng, thường được dùng để biểu thị phân đoạn tối đa: <[], device_ids=[3]>
• Một lưới có hai trục và mã thiết bị ngầm ẩn iota(6): <["a"=2, "b"=3]>
• Lưới có hai trục và mã thiết bị rõ ràng chỉ định thứ tự thiết bị: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Các thông số:

MeshAxisAttr

Trục được đặt tên trong lưới

Cú pháp:

#sdy.mesh_axis<
  ::llvm::StringRef,   # name
  int64_t   # size
>

Các thông số:

OpShardingRuleAttr

Chỉ định cách phân vùng một toán tử.

Cú pháp:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  bool   # is_custom_rule
>

Quy tắc phân đoạn chỉ định cách phân đoạn một toán tử theo nhiều thuộc tính trên toán tử – bất kỳ thuộc tính nào, hình dạng của toán hạng, hình dạng của kết quả, v.v. Ví dụ:

%0 = stablehlo.add %arg0, %arg1 {
    sdy.sharding_rule = #sdy.op_sharding_rule<
        ([i, j],[i, j])->([i, j])
        {i=8, j=8}>
} : tensor<8x8xf32>

%1 = stablehlo.dot_general %arg2, %arg3, contracting_dims = [1] x [0] {
  sdy.sharding_rule = #sdy.op_sharding_rule<
      ([i, k],[k, j])->([i, j])
      {i=8, j=16, k=8}>
}: (tensor<8x8xf32>, tensor<8x16xf32>) -> tensor<8x16xf32>

Lưu ý rằng chúng tôi cho phép các yếu tố có kích thước 1 mặc dù không thể phân đoạn, nhưng việc này chủ yếu là để đảm bảo tính hoàn chỉnh vì nhiều hoạt động như hoạt động trỏ hướng có kích thước một chiều tương ứng với toán hạng và kết quả.

is_custom_rule mô tả liệu đây có phải là quy tắc do người dùng xác định cho thao tác stablehlo.custom_call hay không. Trình phân vùng không biết cách phân vùng các thao tác này, vì vậy, người dùng phải cho biết cách thực hiện. Khi đó là quy tắc tuỳ chỉnh, thì quy tắc đó sẽ luôn được giữ nguyên/không bao giờ bị xoá. is_custom_rule chỉ có thể là true đối với các toán tử stablehlo.custom_call.

Các thông số:

SubAxisInfoAttr

Thông tin về cách trục phụ này được lấy từ trục đầy đủ

Cú pháp:

#sdy.sub_axis_info<
  int64_t,   # pre_size
  int64_t   # size
>

Khi chia một trục đầy đủ thành n trục phụ, trục này được định hình lại thành [k_1,...,k_n] và trục phụ thứ i có thể được biểu thị bằng tích của tất cả kích thước trục ở bên trái m=prod(k_1,...,k_(i-1)) (còn gọi là kích thước trước) và kích thước k_i. Do đó, thuộc tính sub-axis-info chứa hai số đó và được ký hiệu như sau: (m)k cho kích thước trước m và kích thước k.

Các thông số:

TensorMappingAttr

Ánh xạ nhân tố cho mỗi phương diện của một tensor.

Cú pháp:

#sdy.tensor_mapping<
  ::llvm::ArrayRef<DimMappingAttr>   # dim_mappings
>

Các thông số:

TensorShardingAttr

Phân đoạn Tensor

Cú pháp:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>   # replicated_axes
>

Phân đoạn tensor được liên kết với một lưới cụ thể và chỉ có thể tham chiếu tên trục từ lưới đó. Các phân đoạn theo phương diện cho chúng ta biết về từng phương diện của tensor, theo đó các trục (hoặc trục phụ) được phân đoạn từ chính đến phụ. Tất cả các trục khác không phân đoạn một phương diện sẽ được sao chép một cách ngầm ẩn hoặc rõ ràng (nếu các trục đó xuất hiện trong danh sách các trục được sao chép).

Lưới liên kết với tính năng phân đoạn này có thể được chỉ định bằng tên biểu tượng, tham chiếu đến biểu tượng MeshOp tương ứng hoặc MeshAttr cùng dòng.

Các thông số:

TensorShardingPerValueAttr

Phân đoạn Tensor cho mỗi toán hạng/kết quả của một hoạt động

Cú pháp:

#sdy.sharding_per_value<
  ::llvm::ArrayRef<TensorShardingAttr>   # shardings
>

Các thông số:

Enum

PropagationDirection

enum hướng truyền

Trường hợp: