[Mlir-commits] [mlir] 8d05a28 - [mlir][python] Adapt to `segment_sizes` attribute type change.

[Stella Laurenzo]

Author: Stella Laurenzo
Date: 2021-03-19T18:47:00-07:00
New Revision: 8d05a28887ee1e3cbcddf892de8bbc560432afd2

URL: https://github.com/llvm/llvm-project/commit/8d05a28887ee1e3cbcddf892de8bbc560432afd2
DIFF: https://github.com/llvm/llvm-project/commit/8d05a28887ee1e3cbcddf892de8bbc560432afd2.diff

LOG: [mlir][python] Adapt to `segment_sizes` attribute type change.

* Broken by https://reviews.llvm.org/rG1a75be0023cd80fd8560d689999a63d4368c90e6

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/test/Bindings/Python/ods_helpers.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 9d87aa52f7c8..0a4c5fcb40c3 100644
--- 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
index 54f68a82fc01..badeac37034f 100644
--- 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],