Integrate LLVM at llvm/llvm-project@b214ca82daee

Updates LLVM usage to match
[b214ca82daee](llvm/llvm-project@b214ca82daee)

PiperOrigin-RevId: 700689999

Author: d0k

jaxlib/mosaic/dialect/tpu/tpu.td

@@ -44,7 +44,7 @@ class TPU_Attr<string name, string mnemonic_, list<Trait> traits = []>
 }
 
 // TODO(b/369418606): Find out the way to verify vreg size.
-def TPU_Vreg : Type<IsVectorTypePred, "native-sized vreg", "::mlir::VectorType">;
+def TPU_Vreg : Type<IsVectorOfNonZeroRankTypePred, "native-sized vreg", "::mlir::VectorType">;
 
 class TPU_Type<string name, string mnemonic_, list<Trait> traits = []>
     : TypeDef<TPU_Dialect, name, traits> {
@@ -179,8 +179,8 @@ def TPU_ReductionKindAttr
 }
 
 def TPU_AllReduceOp : TPU_Op<"all_reduce", [Pure, SameOperandsAndResultType]> {
-  let arguments = (ins AnyVector:$input, I64Attr:$dim, TPU_ReductionKindAttr:$kind);
-  let results = (outs AnyVector:$output);
+  let arguments = (ins AnyVectorOfNonZeroRank:$input, I64Attr:$dim, TPU_ReductionKindAttr:$kind);
+  let results = (outs AnyVectorOfNonZeroRank:$output);
   let assemblyFormat = [{
     $input attr-dict `:` type($input)
   }];
@@ -217,11 +217,11 @@ def TPU_LoadOp : TPU_Op<"load"> {
 // TODO(jevinjiang): migrate tpu.strided_store to general vector store op.
 def TPU_VectorStoreOp :TPU_Op<"vector_store", [AttrSizedOperandSegments]> {
   let arguments = (ins
-    AnyVector:$valueToStore,
+    AnyVectorOfNonZeroRank:$valueToStore,
     AnyMemRef:$base,
     Variadic<Index>:$indices,
     DenseI32ArrayAttr:$strides,
-    Optional<AnyVector>:$mask   // Elementwise mask.
+    Optional<AnyVectorOfNonZeroRank>:$mask   // Elementwise mask.
   );
   let results = (outs);
   let assemblyFormat = [{
@@ -236,7 +236,7 @@ def TPU_StridedLoadOp : TPU_Op<"strided_load"> {
     Variadic<Index>:$indices,
     DenseI32ArrayAttr:$strides
   );
-  let results = (outs AnyVector:$result);
+  let results = (outs AnyVectorOfNonZeroRank:$result);
   let assemblyFormat = [{
     $base `[` $indices `]` attr-dict `:` type($base) `,` type($result)
   }];
@@ -245,7 +245,7 @@ def TPU_StridedLoadOp : TPU_Op<"strided_load"> {
 
 def TPU_StridedStoreOp : TPU_Op<"strided_store"> {
   let arguments = (ins
-    AnyVector:$valueToStore,
+    AnyVectorOfNonZeroRank:$valueToStore,
     AnyMemRef:$base,
     Variadic<Index>:$indices,
     DenseI32ArrayAttr:$strides
@@ -291,15 +291,15 @@ def TPU_ShuffledStoreOp : TPU_Op<"shuffled_store"> {
 // TODO(jevinjiang): deprecate to use dynamic_rotate.
 def TPU_RotateOp : TPU_Op<"rotate", [Pure, SameOperandsAndResultType]> {
   let arguments = (ins
-    AnyVector:$value,
+    AnyVectorOfNonZeroRank:$value,
     SI32Attr:$amount,
     SI32Attr:$dimension,
     // When the stride is specified, the rotation amount for each index on the
     // stride dimension will be (amount + stride * index).
     OptionalAttr<SI32Attr>:$stride,
     OptionalAttr<SI32Attr>:$stride_dimension
   );
-  let results = (outs AnyVector:$result);
+  let results = (outs AnyVectorOfNonZeroRank:$result);
   let assemblyFormat = [{
     $value `by` $amount `dim` $dimension (`stride` $stride `stride_dim` $stride_dimension^)? attr-dict `:` type($value)
   }];
@@ -308,15 +308,15 @@ def TPU_RotateOp : TPU_Op<"rotate", [Pure, SameOperandsAndResultType]> {
 
 def TPU_DynamicRotateOp : TPU_Op<"dynamic_rotate", [Pure]> {
   let arguments = (ins
-    AnyVector:$value,
+    AnyVectorOfNonZeroRank:$value,
     I32:$amount,
     SI32Attr:$dimension,
     // When the stride is specified, the rotation amount for each index on the
     // stride dimension will be (amount + stride * index).
     OptionalAttr<SI32Attr>:$stride,
     OptionalAttr<SI32Attr>:$stride_dimension
   );
-  let results = (outs AnyVector:$result);
+  let results = (outs AnyVectorOfNonZeroRank:$result);
   let assemblyFormat = [{
     $value `by` $amount `dim` $dimension attr-dict `:` type($value) `,` type($amount)  `->` type($result)
   }];
@@ -325,30 +325,30 @@ def TPU_DynamicRotateOp : TPU_Op<"dynamic_rotate", [Pure]> {
 
 def TPU_IotaOp : TPU_Op<"iota", [Pure]> {
   let arguments = (ins OptionalAttr<I32Attr>:$dimension);
-  let results = (outs AnyVector:$output);
+  let results = (outs AnyVectorOfNonZeroRank:$output);
   let assemblyFormat = [{ attr-dict `:` type($output) }];
 }
 
 // TODO(mvoz): deprecated - use concat. Canonicalization will do so automatically.
 // b/376295711
 def TPU_RepeatOp : TPU_Op<"repeat", [Pure]> {
   let arguments = (ins
-    AnyVector:$source,
+    AnyVectorOfNonZeroRank:$source,
     I32Attr:$dimension,
     I32Attr:$times
   );
-  let results = (outs AnyVector:$output);
+  let results = (outs AnyVectorOfNonZeroRank:$output);
   let assemblyFormat = [{ $source `,` $dimension `x` $times attr-dict `:` type($source) `->` type($output) }];
 }
 
 def TPU_BroadcastInSublanesOp : TPU_Op<"broadcast_in_sublanes", [Pure]> {
   let arguments = (ins
-    AnyVector:$source,  // All sublanes should be equal.
+    AnyVectorOfNonZeroRank:$source,  // All sublanes should be equal.
     I32Attr:$lane  // Coordinates of the first element to take.
   );
   // Output shape should be the same, except for position dim which contains
   // the newly inserted dimension.
-  let results = (outs AnyVector:$output);
+  let results = (outs AnyVectorOfNonZeroRank:$output);
   let assemblyFormat = [{
     $source `,` $lane attr-dict `:` type($source) `->` type($output)
   }];
@@ -357,30 +357,30 @@ def TPU_BroadcastInSublanesOp : TPU_Op<"broadcast_in_sublanes", [Pure]> {
 // Integer unpacks are always signed at the moment.
 def TPU_UnpackSubelementsOp : TPU_Op<"unpack_subelements", [Pure]> {
   let arguments = (ins
-    AnyVector:$source,
+    AnyVectorOfNonZeroRank:$source,
     I32Attr:$index
   );
-  let results = (outs AnyVector:$output);
+  let results = (outs AnyVectorOfNonZeroRank:$output);
   let assemblyFormat = [{ $source `,` $index attr-dict `:` type($source) `->` type($output) }];
 }
 
 // Integer packs are always signed at the moment.
 def TPU_PackSubelementsOp : TPU_Op<"pack_subelements", [Pure]> {
   let arguments = (ins
-    Variadic<AnyVector>:$sources,
+    Variadic<AnyVectorOfNonZeroRank>:$sources,
     TPU_PackFormatEnum:$pack_format
   );
-  let results = (outs AnyVector:$output);
+  let results = (outs AnyVectorOfNonZeroRank:$output);
   let assemblyFormat = [{ $sources attr-dict `:` type($sources) `->` type($output) }];
 }
 
 def TPU_GatherOp : TPU_Op<"gather", [Pure]> {
   let arguments = (ins
-    AnyVector:$source,
+    AnyVectorOfNonZeroRank:$source,
     DenseI32ArrayAttr:$indices,
     I32Attr:$dimension
   );
-  let results = (outs AnyVector:$output);
+  let results = (outs AnyVectorOfNonZeroRank:$output);
   let assemblyFormat = [{
     $source `[` $indices `]` `in` $dimension attr-dict
     `:` type($source) `->` type($output)
@@ -389,11 +389,11 @@ def TPU_GatherOp : TPU_Op<"gather", [Pure]> {
 
 def TPU_DynamicGatherOp : TPU_Op<"dynamic_gather", [Pure]> {
   let arguments = (ins
-    AnyVector:$source,
-    AnyVector:$indices,  // If this is 2D, only the first row matters.
+    AnyVectorOfNonZeroRank:$source,
+    AnyVectorOfNonZeroRank:$indices,  // If this is 2D, only the first row matters.
     I32Attr:$dimension
   );
-  let results = (outs AnyVector:$output);
+  let results = (outs AnyVectorOfNonZeroRank:$output);
   let assemblyFormat = [{
     $source `[` $indices `]` `in` $dimension attr-dict
     `:` type($source) `,` type($indices) `->` type($output)
@@ -424,9 +424,9 @@ def TPU_DotDimensionNumbersAttr : TPU_Attr<"DotDimensionNumbers", "dot_dimension
 // TODO(apaszke): Think hard about precision
 def TPU_MatmulOp : TPU_Op<"matmul", [Pure]> {
   let arguments = (ins
-    AnyVector:$lhs,
-    AnyVector:$rhs,
-    AnyVector:$acc,
+    AnyVectorOfNonZeroRank:$lhs,
+    AnyVectorOfNonZeroRank:$rhs,
+    AnyVectorOfNonZeroRank:$acc,
     // These flags are deprecated - if dimension_numbers are defined,
     // these flags are ignored. They will always be false after canonicalize.
     DefaultValuedAttr<BoolAttr, "false">:$transpose_lhs,
@@ -435,7 +435,7 @@ def TPU_MatmulOp : TPU_Op<"matmul", [Pure]> {
     // NOTE: User-level optional, once canonicalized, always present.
     OptionalAttr<TPU_DotDimensionNumbersAttr>:$dimension_numbers
   );
-  let results = (outs AnyVector:$result);
+  let results = (outs AnyVectorOfNonZeroRank:$result);
   let assemblyFormat = [{
     $lhs `,` $rhs `,` $acc attr-dict `:` type($lhs) `,` type($rhs) `,` type($acc) `->` type($result)
   }];
@@ -445,19 +445,19 @@ def TPU_MatmulOp : TPU_Op<"matmul", [Pure]> {
 
 def TPU_ConcatenateOp : TPU_Op<"concatenate", [Pure]> {
   let arguments = (ins
-    Variadic<AnyVector>:$sources,
+    Variadic<AnyVectorOfNonZeroRank>:$sources,
     I32Attr:$dimension
   );
-  let results = (outs AnyVector:$output);
+  let results = (outs AnyVectorOfNonZeroRank:$output);
   let assemblyFormat = [{
     $sources `in` $dimension attr-dict `:` type($sources) `->` type($output)
   }];
   let hasVerifier = 1;
 }
 
 def TPU_BitcastOp : TPU_Op<"bitcast", [Pure]> {
-  let arguments = (ins AnyVector:$input);
-  let results = (outs AnyVector:$output);
+  let arguments = (ins AnyVectorOfNonZeroRank:$input);
+  let results = (outs AnyVectorOfNonZeroRank:$output);
   let assemblyFormat = [{ $input attr-dict `:` type($input) `->` type($output) }];
   let hasVerifier = 1;
 }
@@ -469,16 +469,16 @@ def TPU_BitcastVregOp : TPU_Op<"bitcast_vreg", [Pure]> {
 }
 
 def TPU_RollVectorsOp : TPU_Op<"roll_vectors", [Pure]> {
-  let arguments = (ins Variadic<AnyVector>:$input);
-  let results = (outs AnyVector:$output);
+  let arguments = (ins Variadic<AnyVectorOfNonZeroRank>:$input);
+  let results = (outs AnyVectorOfNonZeroRank:$output);
   let assemblyFormat = [{
     $input attr-dict `:` type($input) `->` type($output)
   }];
 }
 
 def TPU_UnrollVectorsOp : TPU_Op<"unroll_vectors", [Pure]> {
-  let arguments = (ins AnyVector:$input);
-  let results = (outs Variadic<AnyVector>:$output);
+  let arguments = (ins AnyVectorOfNonZeroRank:$input);
+  let results = (outs Variadic<AnyVectorOfNonZeroRank>:$output);
   let hasCanonicalizeMethod = 1;
   let assemblyFormat = [{
     $input attr-dict `:` type($input) `->` type($output)
@@ -722,8 +722,8 @@ def TPU_DelayOp : TPU_Op<"delay"> {
 
 // Expands the granularity of mask to subelements.
 def TPU_MaskCastOp : TPU_Op<"mask_cast", [Pure]> {
-  let arguments = (ins AnyVector:$input);
-  let results = (outs AnyVector:$result);
+  let arguments = (ins AnyVectorOfNonZeroRank:$input);
+  let results = (outs AnyVectorOfNonZeroRank:$result);
   let assemblyFormat = [{
     $input attr-dict `:` type($input) `->` type($result)
   }];
@@ -749,7 +749,7 @@ def TPU_PRNGSeed32Op : TPU_Op<"prng_set_seed_32"> {
 
 def TPU_PRNGRandomBitsOp : TPU_Op<"prng_random_bits"> {
   let arguments = (ins);
-  let results = (outs AnyVector:$output);
+  let results = (outs AnyVectorOfNonZeroRank:$output);
 }
 
 def TPU_LogOp : TPU_Op<"log"> {