use size_t

Signed-off-by: Bangtian Liu <[email protected]>

Author: bangtianliu

mlir/include/mlir-c/Dialect/Linalg.h

@@ -37,7 +37,7 @@ mlirLinalgInferContractionDimensions(MlirOperation op);
 
 MLIR_CAPI_EXPORTED MlirLinalgContractionDimensions
 mlirLinalgInferContractionDimensionsFromMaps(const MlirAffineMap *indexingMaps,
-                                             intptr_t numMaps);
+                                             size_t numMaps);
 
 MLIR_CAPI_EXPORTED bool mlirLinalgIsAConvolutionOp(MlirOperation op);
 

mlir/lib/CAPI/Dialect/Linalg.cpp

@@ -77,14 +77,14 @@ mlirLinalgInferContractionDimensions(MlirOperation op) {
 
 MLIR_CAPI_EXPORTED MlirLinalgContractionDimensions
 mlirLinalgInferContractionDimensionsFromMaps(const MlirAffineMap *indexingMaps,
-                                             intptr_t numMaps) {
+                                             size_t numMaps) {
   MlirLinalgContractionDimensions result{};
-  if (!indexingMaps || numMaps <= 0)
+  if (!indexingMaps || numMaps == 0)
     return result;
 
-  SmallVector<AffineMap> maps;
+  SmallVector<AffineMap, 3> maps;
   maps.reserve(numMaps);
-  for (intptr_t i = 0; i < numMaps; ++i) {
+  for (size_t i = 0; i < numMaps; ++i) {
     maps.push_back(unwrap(indexingMaps[i]));
   }
 

mlir/test/python/dialects/linalg/utils.py

@@ -237,15 +237,14 @@ def test_infer_contraction_dimensions_from_maps():
             assert invalid_dims is None
 
             # === Test element-wise operation ===
-            # All dimensions appear in all operands, so they're batch dimensions.
             dim_i = AffineDimExpr.get(0)
             dim_j = AffineDimExpr.get(1)
             elementwise_map = AffineMap.get(2, 0, [dim_i, dim_j])
             elementwise_dims = linalg.infer_contraction_dimensions_from_maps(
                 [elementwise_map, elementwise_map, elementwise_map]
             )
             assert elementwise_dims is not None
-            assert list(elementwise_dims.m) == []
-            assert list(elementwise_dims.n) == []
-            assert list(elementwise_dims.k) == []
+            assert len(elementwise_dims.m) == 0
+            assert len(elementwise_dims.n) == 0
+            assert len(elementwise_dims.k) == 0
             assert list(elementwise_dims.batch) == [0, 1]