[mlir] Add isStatic* size check for ShapedTypes. NFCI. (#147085)

The motivation is to avoid having to negate `isDynamic*` checks, avoid
double negations, and allow for `ShapedType::isStaticDim` to be used in
ADT functions without having to wrap it in a lambda performing the
negation.

Also add the new functions to C and Python bindings.

Author: kuhar

mlir/include/mlir-c/BuiltinTypes.h

@@ -289,9 +289,12 @@ MLIR_CAPI_EXPORTED int64_t mlirShapedTypeGetRank(MlirType type);
 /// Checks whether the given shaped type has a static shape.
 MLIR_CAPI_EXPORTED bool mlirShapedTypeHasStaticShape(MlirType type);
 
-/// Checks wither the dim-th dimension of the given shaped type is dynamic.
+/// Checks whether the dim-th dimension of the given shaped type is dynamic.
 MLIR_CAPI_EXPORTED bool mlirShapedTypeIsDynamicDim(MlirType type, intptr_t dim);
 
+/// Checks whether the dim-th dimension of the given shaped type is static.
+MLIR_CAPI_EXPORTED bool mlirShapedTypeIsStaticDim(MlirType type, intptr_t dim);
+
 /// Returns the dim-th dimension of the given ranked shaped type.
 MLIR_CAPI_EXPORTED int64_t mlirShapedTypeGetDimSize(MlirType type,
                                                     intptr_t dim);
@@ -300,17 +303,25 @@ MLIR_CAPI_EXPORTED int64_t mlirShapedTypeGetDimSize(MlirType type,
 /// in shaped types.
 MLIR_CAPI_EXPORTED bool mlirShapedTypeIsDynamicSize(int64_t size);
 
+/// Checks whether the given shaped type dimension value is statically-sized.
+MLIR_CAPI_EXPORTED bool mlirShapedTypeIsStaticSize(int64_t size);
+
 /// Returns the value indicating a dynamic size in a shaped type. Prefer
-/// mlirShapedTypeIsDynamicSize to direct comparisons with this value.
+/// mlirShapedTypeIsDynamicSize and mlirShapedTypeIsStaticSize to direct
+/// comparisons with this value.
 MLIR_CAPI_EXPORTED int64_t mlirShapedTypeGetDynamicSize(void);
 
 /// Checks whether the given value is used as a placeholder for dynamic strides
 /// and offsets in shaped types.
 MLIR_CAPI_EXPORTED bool mlirShapedTypeIsDynamicStrideOrOffset(int64_t val);
 
+/// Checks whether the given dimension value of a stride or an offset is
+/// statically-sized.
+MLIR_CAPI_EXPORTED bool mlirShapedTypeIsStaticStrideOrOffset(int64_t val);
+
 /// Returns the value indicating a dynamic stride or offset in a shaped type.
-/// Prefer mlirShapedTypeGetDynamicStrideOrOffset to direct comparisons with
-/// this value.
+/// Prefer mlirShapedTypeIsDynamicStrideOrOffset and
+/// mlirShapedTypeIsStaticStrideOrOffset to direct comparisons with this value.
 MLIR_CAPI_EXPORTED int64_t mlirShapedTypeGetDynamicStrideOrOffset(void);
 
 //===----------------------------------------------------------------------===//

mlir/lib/Bindings/Python/IRTypes.cpp

@@ -544,6 +544,15 @@ void mlir::PyShapedType::bindDerived(ClassTy &c) {
       nb::arg("dim"),
       "Returns whether the dim-th dimension of the given shaped type is "
       "dynamic.");
+  c.def(
+      "is_static_dim",
+      [](PyShapedType &self, intptr_t dim) -> bool {
+        self.requireHasRank();
+        return mlirShapedTypeIsStaticDim(self, dim);
+      },
+      nb::arg("dim"),
+      "Returns whether the dim-th dimension of the given shaped type is "
+      "static.");
   c.def(
       "get_dim_size",
       [](PyShapedType &self, intptr_t dim) {
@@ -558,6 +567,12 @@ void mlir::PyShapedType::bindDerived(ClassTy &c) {
       nb::arg("dim_size"),
       "Returns whether the given dimension size indicates a dynamic "
       "dimension.");
+  c.def_static(
+      "is_static_size",
+      [](int64_t size) -> bool { return mlirShapedTypeIsStaticSize(size); },
+      nb::arg("dim_size"),
+      "Returns whether the given dimension size indicates a static "
+      "dimension.");
   c.def(
       "is_dynamic_stride_or_offset",
       [](PyShapedType &self, int64_t val) -> bool {
@@ -567,6 +582,15 @@ void mlir::PyShapedType::bindDerived(ClassTy &c) {
       nb::arg("dim_size"),
       "Returns whether the given value is used as a placeholder for dynamic "
       "strides and offsets in shaped types.");
+  c.def(
+      "is_static_stride_or_offset",
+      [](PyShapedType &self, int64_t val) -> bool {
+        self.requireHasRank();
+        return mlirShapedTypeIsStaticStrideOrOffset(val);
+      },
+      nb::arg("dim_size"),
+      "Returns whether the given shaped type stride or offset value is "
+      "statically-sized.");
   c.def_prop_ro(
       "shape",
       [](PyShapedType &self) {

mlir/lib/CAPI/IR/BuiltinTypes.cpp

@@ -332,6 +332,11 @@ bool mlirShapedTypeIsDynamicDim(MlirType type, intptr_t dim) {
       .isDynamicDim(static_cast<unsigned>(dim));
 }
 
+bool mlirShapedTypeIsStaticDim(MlirType type, intptr_t dim) {
+  return llvm::cast<ShapedType>(unwrap(type))
+      .isStaticDim(static_cast<unsigned>(dim));
+}
+
 int64_t mlirShapedTypeGetDimSize(MlirType type, intptr_t dim) {
   return llvm::cast<ShapedType>(unwrap(type))
       .getDimSize(static_cast<unsigned>(dim));
@@ -343,10 +348,18 @@ bool mlirShapedTypeIsDynamicSize(int64_t size) {
   return ShapedType::isDynamic(size);
 }
 
+bool mlirShapedTypeIsStaticSize(int64_t size) {
+  return ShapedType::isStatic(size);
+}
+
 bool mlirShapedTypeIsDynamicStrideOrOffset(int64_t val) {
   return ShapedType::isDynamic(val);
 }
 
+bool mlirShapedTypeIsStaticStrideOrOffset(int64_t val) {
+  return ShapedType::isStatic(val);
+}
+
 int64_t mlirShapedTypeGetDynamicStrideOrOffset() {
   return ShapedType::kDynamic;
 }

mlir/python/mlir/_mlir_libs/_mlir/ir.pyi

@@ -2497,6 +2497,11 @@ class ShapedType(Type):
         Returns whether the given dimension size indicates a dynamic dimension.
         """
     @staticmethod
+    def is_static_size(dim_size: int) -> bool:
+        """
+        Returns whether the given dimension size indicates a static dimension.
+        """
+    @staticmethod
     def isinstance(other: Type) -> bool: ...
     def __init__(self, cast_from_type: Type) -> None: ...
     def get_dim_size(self, dim: int) -> int:
@@ -2507,10 +2512,18 @@ class ShapedType(Type):
         """
         Returns whether the dim-th dimension of the given shaped type is dynamic.
         """
+    def is_static_dim(self, dim: int) -> bool:
+        """
+        Returns whether the dim-th dimension of the given shaped type is static.
+        """
     def is_dynamic_stride_or_offset(self, dim_size: int) -> bool:
         """
         Returns whether the given value is used as a placeholder for dynamic strides and offsets in shaped types.
         """
+    def is_static_stride_or_offset(self, dim_size: int) -> bool:
+        """
+        Returns whether the given shaped type stride or offset value is statically-sized.
+        """
     @property
     def element_type(self) -> Type:
         """