Integrate LLVM at llvm/llvm-project@fc44a4fcd3c5 (#2851)

Author: mrguenther

WORKSPACE.bazel

@@ -17,9 +17,9 @@ workspace(name = "stablehlo")
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
-LLVM_COMMIT = "6d4a0935c850ec3ddfc70c4ba97b98adc35c676e"
+LLVM_COMMIT = "fc44a4fcd3c54be927c15ddd9211aca1501633e7"
 
-LLVM_SHA256 = "01d73796d7c614c809fe464b4beb1cfb5e275e805e7b985ef90c97fe22f01154"
+LLVM_SHA256 = "d228aebe5583c69c4e48fd7a8e149e3d22ee6dafaeae94009467143d32d9bfc4"
 
 http_archive(
     name = "llvm-raw",

build_tools/github_actions/lint_version.sh

@@ -43,7 +43,7 @@ VERSIONED_COMPAT_TEST="$COMPAT_TEST_BASE.$TEST_VERSION.mlir"
 VERSIONED_COMPAT_TEST_BC="$COMPAT_TEST_BASE.$TEST_VERSION.mlir.bc"
 
 show_help() {
-  HELP_URL="https://github.com/openxla/stablehlo/blob/main/docs/vhlo.md#add-versioned-serialization-test"
+  HELP_URL="https://github.com/openxla/stablehlo/blob/main/docs/vhlo_checklist.md#4-add-versioned-serialization-test"
   echo "For details on creating versioned tests for a new minor version of"
   echo "StableHLO, see the instructions on:"
   echo "$HELP_URL"

build_tools/llvm_version.txt

@@ -1 +1 @@
-6d4a0935c850ec3ddfc70c4ba97b98adc35c676e
+fc44a4fcd3c54be927c15ddd9211aca1501633e7

docs/spec.md

@@ -89,7 +89,7 @@ completely numeric to simplify generation of StableHLO programs.
 
 ```ebnf
 Type         ::= ValueType | NonValueType
-ValueType    ::= TensorType | QuantizedTensorType | TokenType | TupleType
+ValueType    ::= TensorType | QuantizedTensorType | TokenType | TupleType | BufferType
 NonValueType ::= TensorElementType | QuantizedTensorElementType | FunctionType | StringType
 ```
 
@@ -229,6 +229,21 @@ TupleType ::= 'tuple' '<' TupleElementTypes '>'
 TupleElementTypes ::= [ValueType {',' ValueType}]
 ```
 
+**Buffer types** represent buffers. For example, in XLA, buffers are
+multidimensional arrays with consistent storage. Similar to **tensor types**,
+buffer types have a **shape** and an **element type**, where a shape represents
+non-negative or unknown **dimension sizes** in the ascending order of the
+corresponding **dimensions** (which are also called **axes**) numbered from `0`
+to `R-1`. The number of dimensions `R` is called **rank**. For example,
+`memref<2x3xf32>` is a buffer type with shape `2x3` and element type `f32`. It
+has two dimensions (or, in other words, two axes) - 0th dimension and 1st
+dimension - whose sizes are 2 and 3. Its rank is 2.
+
+Buffers can be allocated using a `custom_call` to `CreateBuffer` or `Pin` and
+deallocated via a `custom_call` to `Unpin`. Only `custom_call` ops can read and
+write the content inside buffers. See [custom_call](#custom_call) for more
+detail.
+
 **Tuple types** represent tuples, i.e. heterogeneous lists. Tuples are a legacy
 feature which only exists for compatibility with HLO. In HLO, tuples are
 used to represent variadic inputs and outputs. In StableHLO, variadic inputs and
@@ -2433,21 +2448,63 @@ the XLA compiler. In the future, we are planning to unify this metadata
 
 #### Inputs
 
-| Label | Name                  | Type                                              |
-|-------|-----------------------|---------------------------------------------------|
-| (I1)  | `inputs`              | variadic number of values                         |
-| (I2)  | `call_target_name`    | constant of type `string`                         |
-| (I3)  | `has_side_effect`     | constant of type `i1`                             |
-| (I4)  | `backend_config`      | constant of type `string` or attribute dictionary |
-| (I5)  | `api_version`         | constant of type `si32`                           |
-| (I6)  | `called_computations` | variadic number of constants of type `string`     |
+| Label | Name                     | Type                                                       |
+|-------|--------------------------|------------------------------------------------------------|
+| (I1)  | `inputs`                 | variadic number of values                                  |
+| (I2)  | `call_target_name`       | constant of type `string`                                  |
+| (I3)  | `has_side_effect`        | constant of type `i1`                                      |
+| (I4)  | `backend_config`         | constant of type `string` or attribute dictionary          |
+| (I5)  | `api_version`            | constant of type `si32`                                    |
+| (I6)  | `called_computations`    | variadic number of constants of type `string`              |
+| (I7)  | `output_operand_aliases` | specify the aliasing parts in the outputs and operands     |
 
 #### Outputs
 
 | Name      | Type                      |
 |-----------|---------------------------|
 | `results` | variadic number of values |
 
+### (XLA GPU Support) Special custom_call targets
+
+There are three special `call_target_name` related to `buffer` types:
+`CreateBuffer` creates an uninitialized `buffer`, `Pin` creates an initialized
+`buffer` and `Unpin` deallocates a `buffer` and returns the content of the
+`buffer`.
+
+```mlir
+%uninitialized_buffer = "stablehlo.custom_call"() {
+  call_target_name = "CreateBuffer",
+  api_version = 4 : i32,
+} : () -> memref<4xf64>
+
+%initialized_buffer = "stablehlo.custom_call"(%init_value) {
+  call_target_name = "Pin",
+  api_version = 4 : i32,
+} : (tensor<4xf64>) -> memref<4xf64>
+
+%dealloc_buffer = "stablehlo.custom_call"(%initialized_buffer) {
+  call_target_name = "Unpin",
+  api_version = 4 : i32,
+} : (memref<4xf64>) -> tensor<4xf64>
+
+```
+
+### Alias
+
+Some custom_call ops may require a part in the outputs and a part in the
+operands to share the same memory. This can be expressed via
+`output_operand_aliases`. An alias pair representation consists a list of output
+ tuple indices representing the output part, and an operand_index along with a
+ list of operand tuple indices representing the operand part. The list of output
+  or operand tuple indices is empty if the corresponding type is not a `tuple`
+  type, and can be arbitrarily long for an arbitrarily nested tuple type. This
+  is similar to [the XLA alias representation](https://www.tensorflow.org/xla/aliasing).
+
+The output part and the input part in an alias pair must have the same type. For
+custom_call ops that aren't call to `CreateBuffer`, `Pin` and `Unpin`, a
+`buffer` operand can appear in at most one pair of alias, and a `buffer` output
+must appear in one pair of alias.
+
 #### Examples
 
 ```mlir
@@ -2458,6 +2515,16 @@ the XLA compiler. In the future, we are planning to unify this metadata
   api_version = 4 : i32,
   called_computations = [@foo]
 } : (tensor<f64>) -> tensor<f64>
+
+%updated_buffer = "stablehlo.custom_call"(%buffer) {
+  call_target_name = "Update",
+  api_version = 4 : i32,
+  output_operand_aliases = [
+    #stablehlo.output_operand_alias<output_tuple_indices = [],
+      operand_index = 0,
+      operand_tuple_indices = []>]
+} : (memref<4xf64>) -> memref<4xf64>
+
 ```
 
 ### divide
@@ -3780,9 +3847,9 @@ Extracts element at `index` position of the `operand` tuple and produces a
 
 #### Outputs
 
-| Name     | Type               | Constraints |
-|----------|--------------------|-------------|
-| `result` | any supported type | (C2)        |
+| Name     | Type                   | Constraints |
+|----------|------------------------|-------------|
+| `result` | any value              | (C2)        |
 
 #### Constraints
 
@@ -6583,10 +6650,10 @@ Produces a `result` tuple from values `val`.
 #### Examples
 
 ```mlir
-// %val0: [1.0, 2.0]
+// %val0: memref[1.0, 2.0]
 // %val1: (3)
-%result = "stablehlo.tuple"(%val0, %val1) : (tensor<2xf32>, tuple<tensor<i32>>) -> tuple<tensor<2xf32>, tuple<tensor<i32>>>
-// %result: ([1.0, 2.0], (3))
+%result = "stablehlo.tuple"(%val0, %val1) : (memref<2xf32>, tuple<tensor<i32>>) -> tuple<memref<2xf32>, tuple<tensor<i32>>>
+// %result: (memref[1.0, 2.0], (3))
 ```
 
 &nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/tuple_and_get_tuple_element.mlir)
@@ -6692,17 +6759,17 @@ The behavior of an infinite loop is TBD
 
 #### Inputs
 
-| Label | Name      | Type                                                    | Constraints |
-|-------|-----------|---------------------------------------------------------|-------------|
-| (I1)  | `operand` | variadic number of tensors, quantized tensors or tokens | (C1-C3)     |
-| (I2)  | `cond`    | function                                                | (C1)        |
-| (I3)  | `body`    | function                                                | (C2)        |
+| Label | Name      | Type                                    | Constraints |
+|-------|-----------|-----------------------------------------|-------------|
+| (I1)  | `operand` | variadic number of values               | (C1-C3)     |
+| (I2)  | `cond`    | function                                | (C1)        |
+| (I3)  | `body`    | function                                | (C2)        |
 
 #### Outputs
 
-| Name      | Type                                                    | Constraints |
-|-----------|---------------------------------------------------------|-------------|
-| `results` | variadic number of tensors, quantized tensors or tokens | (C3)        |
+| Name      | Type                                            | Constraints |
+|-----------|-------------------------------------------------|-------------|
+| `results` | variadic number of values                       | (C3)        |
 
 #### Constraints
 

stablehlo/dialect/Base.cpp

@@ -798,5 +798,35 @@ bool hasSingleBoundedDimension(Type type) {
   return numBoundedDims == 1 && numDynamicDims == 1;
 }
 
+//===----------------------------------------------------------------------===//
+// Utils for type traversal.
+//===----------------------------------------------------------------------===//
+
+namespace {
+LogicalResult mapOverLeafTypesImpl(
+    Type type, function_ref<LogicalResult(Type type, ArrayRef<int64_t>)> fn,
+    std::vector<int64_t>& indices) {
+  if (!isa<TupleType>(type)) {
+    return fn(type, indices);
+  }
+
+  auto tupleType = cast<TupleType>(type);
+  for (size_t i = 0; i < tupleType.size(); ++i) {
+    indices.push_back(i);
+    if (failed(mapOverLeafTypesImpl(tupleType.getType(i), fn, indices)))
+      return failure();
+    indices.pop_back();
+  }
+
+  return success();
+}
+}  // namespace
+
+LogicalResult mapOverLeafTypes(
+    Type type, function_ref<LogicalResult(Type, ArrayRef<int64_t>)> fn) {
+  std::vector<int64_t> indices;
+  return mapOverLeafTypesImpl(type, fn, indices);
+}
+
 }  // namespace hlo
 }  // namespace mlir

stablehlo/dialect/Base.h

@@ -314,6 +314,11 @@ void writeEnumAttribute(EnumTypeAttr val, DialectBytecodeWriter &writer) {
 mlir::Speculation::Speculatability getShapedSpeculatability(Operation *op,
                                                             int64_t shapeCount);
 
+// Applies `fn` to `type` if it is not a `tuple` type. Otherwise, applies `fn`
+// to each leaf type in the `tuple` type tree or until a `fn` returns failure.
+LogicalResult mapOverLeafTypes(
+    Type type, function_ref<LogicalResult(Type, ArrayRef<int64_t>)> fn);
+
 namespace OpTrait {
 
 template <typename ConcreteType>

stablehlo/dialect/Base.td

@@ -193,16 +193,22 @@ def HLO_FloatOrQuantizedIntOrPerAxisQuantizedIntTensor : RankedTensorOf<[HLO_Flo
 
 def HLO_ComplexTensor : RankedTensorOf<[HLO_Complex]>;
 
-def HLO_Tuple : NestedTupleOf<[HLO_Tensor, HLO_PerAxisQuantizedIntTensor, HLO_Token]>;
+def HLO_Buffer : MemRefOf<[HLO_Float, HLO_Pred, HLO_Int, HLO_Complex, HLO_QuantizedInt]>;
+
+def HLO_Tuple : NestedTupleOf<[HLO_Tensor, HLO_Buffer, HLO_PerAxisQuantizedIntTensor, HLO_Token]>;
 
 def HLO_TensorOrToken : AnyTypeOf<[HLO_Tensor, HLO_Token]>;
 
 def HLO_TensorOrPerAxisQuantizedTensorOrToken : AnyTypeOf<[HLO_Tensor, HLO_PerAxisQuantizedIntTensor, HLO_Token]>;
 
+def HLO_TensorOrPerAxisQuantizedTensorOrTokenOrBuffer : AnyTypeOf<[HLO_TensorOrPerAxisQuantizedTensorOrToken, HLO_Buffer]>;
+
 def HLO_TensorOrTokenOrTuple : AnyTypeOf<[HLO_Tensor, HLO_Token, HLO_Tuple]>;
 
 def HLO_TensorOrPerAxisQuantizedTensorOrTokenOrTuple : AnyTypeOf<[HLO_Tensor, HLO_PerAxisQuantizedIntTensor, HLO_Token, HLO_Tuple]>;
 
+def HLO_TensorOrPerAxisQuantizedTensorOrTokenOrTupleOrBuffer : AnyTypeOf<[HLO_TensorOrPerAxisQuantizedTensorOrTokenOrTuple, HLO_Buffer]>;
+
 def HLO_DimensionValue : AnyTypeOf<[Index, HLO_Int]>;
 
 // Dynamic representation of a shape vector as a tensor.
@@ -227,9 +233,9 @@ def HLO_AnyFpOrComplexTensor : TensorOf<[HLO_Float, HLO_Complex]>;
 
 def HLO_AnyPredOrIntTensor : TensorOf<[HLO_Pred, HLO_Int]>;
 
-def HLO_AnyTuple : NestedTupleOf<[HLO_AnyTensor, HLO_Token]>;
+def HLO_CustomCallTuple : NestedTupleOf<[HLO_AnyTensor, HLO_Buffer, HLO_Token]>;
 
-def HLO_CustomCallValue : AnyTypeOf<[HLO_AnyTensor, HLO_Token, HLO_AnyTuple]>;
+def HLO_CustomCallValue : AnyTypeOf<[HLO_AnyTensor, HLO_Buffer, HLO_Token, HLO_CustomCallTuple]>;
 
 //===----------------------------------------------------------------------===//
 // HLO combined type definitions.