[mlir][python] Adapt to segment_sizes attribute type change.

* Broken by https://reviews.llvm.org/rG1a75be0023cd80fd8560d689999a63d4368c90e6
llvm · Mar 20, 2021 · 8d05a28 · 8d05a28
1 parent f380066
commit 8d05a28
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -1034,8 +1034,8 @@ PyOpView::buildGeneric(py::object cls, py::list resultTypeList,
   py::object operandSegmentSpecObj = cls.attr("_ODS_OPERAND_SEGMENTS");
   py::object resultSegmentSpecObj = cls.attr("_ODS_RESULT_SEGMENTS");
 
-  std::vector<uint64_t> operandSegmentLengths;
-  std::vector<uint64_t> resultSegmentLengths;
+  std::vector<uint32_t> operandSegmentLengths;
+  std::vector<uint32_t> resultSegmentLengths;
 
   // Validate/determine region count.
   auto opRegionSpec = py::cast<std::tuple<int, bool>>(cls.attr("_ODS_REGIONS"));
@@ -1247,8 +1247,8 @@ PyOpView::buildGeneric(py::object cls, py::list resultTypeList,
     // Add result_segment_sizes attribute.
     if (!resultSegmentLengths.empty()) {
       int64_t size = resultSegmentLengths.size();
-      MlirAttribute segmentLengthAttr = mlirDenseElementsAttrUInt64Get(
-          mlirVectorTypeGet(1, &size, mlirIntegerTypeGet(context->get(), 64)),
+      MlirAttribute segmentLengthAttr = mlirDenseElementsAttrUInt32Get(
+          mlirVectorTypeGet(1, &size, mlirIntegerTypeGet(context->get(), 32)),
           resultSegmentLengths.size(), resultSegmentLengths.data());
       (*attributes)["result_segment_sizes"] =
           PyAttribute(context, segmentLengthAttr);
@@ -1257,8 +1257,8 @@ PyOpView::buildGeneric(py::object cls, py::list resultTypeList,
     // Add operand_segment_sizes attribute.
     if (!operandSegmentLengths.empty()) {
       int64_t size = operandSegmentLengths.size();
-      MlirAttribute segmentLengthAttr = mlirDenseElementsAttrUInt64Get(
-          mlirVectorTypeGet(1, &size, mlirIntegerTypeGet(context->get(), 64)),
+      MlirAttribute segmentLengthAttr = mlirDenseElementsAttrUInt32Get(
+          mlirVectorTypeGet(1, &size, mlirIntegerTypeGet(context->get(), 32)),
           operandSegmentLengths.size(), operandSegmentLengths.data());
       (*attributes)["operand_segment_sizes"] =
           PyAttribute(context, segmentLengthAttr);

diff --git a/mlir/test/Bindings/Python/ods_helpers.py b/mlir/test/Bindings/Python/ods_helpers.py
@@ -125,8 +125,8 @@ class TestOp(OpView):
       # CHECK: %[[V2:.+]] = "custom.value"
       # CHECK: %[[V3:.+]] = "custom.value"
       # CHECK: "custom.test_op"(%[[V0]], %[[V1]], %[[V2]], %[[V3]])
-      # CHECK-SAME: operand_segment_sizes = dense<[1, 2, 1]> : vector<3xi64>
-      # CHECK-SAME: result_segment_sizes = dense<[2, 1, 1]> : vector<3xi64>
+      # CHECK-SAME: operand_segment_sizes = dense<[1, 2, 1]> : vector<3xi32>
+      # CHECK-SAME: result_segment_sizes = dense<[2, 1, 1]> : vector<3xi32>
       # CHECK-SAME: : (i32, i32, i32, i32) -> (i8, i16, i32, i64)
       op = TestOp.build_generic(
           results=[[t0, t1], t2, t3],