[python] Properly type arange result (#205)

Author: Hardcode84

numba_dpcomp/numba_dpcomp/mlir/numpy/funcs.py

@@ -396,6 +396,16 @@ def ones_like_impl(builder, arr):
     return _init_impl(builder, arr.shape, arr.dtype, 1)
 
 
+_is_np_long64 = numpy.int_ == numpy.int64
+
+
+def _get_numpy_long(builder):
+    if _is_np_long64:
+        return builder.int64
+    else:
+        return builder.int32
+
+
 @register_func("numpy.arange", numpy.arange)
 def arange_impl(builder, start, stop=None, step=None, dtype=None):
     if stop is None:
@@ -406,7 +416,7 @@ def arange_impl(builder, start, stop=None, step=None, dtype=None):
         step = 1
 
     if dtype is None:
-        dtype = builder.int64
+        dtype = _get_numpy_long(builder)
 
     inc = builder.select(step < 0, 1, -1)
     count = (stop - start + step + inc) // step

numba_dpcomp/numba_dpcomp/mlir/tests/test_numba_parfor.py

@@ -107,10 +107,8 @@ def _gen_tests():
         "test_parfor_array_access4",  # np.dot unsupported args
         "test_tuple_concat_with_reverse_slice",  # enumerate
         "test_reduce",  # functools.reduce
-        "test_two_d_array_reduction",  # 'memref<?x?xi64>' and result type 'memref<?x?xi32>' are cast incompatible
         "test_tuple_concat",  # enumerate
         "test_two_d_array_reduction_with_float_sizes",  # np.array
-        "test_two_d_array_reduction_reuse",  # 'memref<?x?xi64>' and result type 'memref<?x?xi32>' are cast incompatible
         "test_parfor_array_access_lower_slice",  # plier.getitem
         "test_size_assertion",  # AssertionError not raised
         "test_parfor_slice18",  # cast types mismatch

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/plier_to_linalg.cpp

@@ -1616,6 +1616,67 @@ struct LowerTupleCasts : public mlir::OpConversionPattern<plier::CastOp> {
   }
 };
 
+static mlir::Value convertTensorElements(mlir::OpBuilder &builder,
+                                         mlir::Location loc, mlir::Value src,
+                                         mlir::ShapedType dstType,
+                                         mlir::Type origSrcElemType,
+                                         mlir::Type origDstElemType) {
+  assert(src.getType().isa<mlir::ShapedType>());
+  auto srcType = src.getType().cast<mlir::ShapedType>();
+  assert(srcType.getRank() == dstType.getRank());
+  if (srcType.getElementType() == dstType.getElementType())
+    return src;
+
+  if (srcType.isa<mlir::MemRefType>())
+    src = builder.create<mlir::bufferization::ToTensorOp>(loc, src);
+
+  auto rank = static_cast<unsigned>(srcType.getRank());
+  llvm::SmallVector<mlir::Value> shape(rank);
+  for (auto i : llvm::seq(0u, rank))
+    shape[i] = builder.create<mlir::tensor::DimOp>(loc, src, i);
+
+  mlir::Value init = builder.create<mlir::linalg::InitTensorOp>(
+      loc, shape, dstType.getElementType());
+
+  auto affineMap =
+      mlir::AffineMap::getMultiDimIdentityMap(rank, builder.getContext());
+  const mlir::AffineMap maps[] = {
+      affineMap,
+      affineMap,
+  };
+
+  llvm::SmallVector<mlir::StringRef> iterators(rank, "parallel");
+
+  auto bodyBuilder = [&](mlir::OpBuilder &b, mlir::Location l,
+                         mlir::ValueRange args) {
+    assert(args.size() == 2);
+    auto doSignCast = [&](mlir::Value src, mlir::Type dstType) -> mlir::Value {
+      if (src.getType() != dstType)
+        return b.create<plier::SignCastOp>(l, dstType, src);
+
+      return src;
+    };
+    auto arg = doSignCast(args.front(), origSrcElemType);
+    arg = b.create<plier::CastOp>(l, origDstElemType, arg);
+    arg = doSignCast(arg, dstType.getElementType());
+    b.create<mlir::linalg::YieldOp>(l, arg);
+  };
+  mlir::Value res =
+      builder
+          .create<mlir::linalg::GenericOp>(loc, init.getType(), src, init, maps,
+                                           iterators, bodyBuilder)
+          .getResult(0);
+
+  if (dstType.isa<mlir::MemRefType>()) {
+    auto memrefType =
+        mlir::MemRefType::get(dstType.getShape(), dstType.getElementType());
+    res = builder.create<mlir::bufferization::ToMemrefOp>(loc, memrefType, res);
+  }
+
+  rerunScfPipeline(res.getDefiningOp());
+  return res;
+}
+
 struct LowerTensorCasts : public mlir::OpConversionPattern<plier::CastOp> {
   using OpConversionPattern::OpConversionPattern;
 
@@ -1653,6 +1714,9 @@ struct LowerTensorCasts : public mlir::OpConversionPattern<plier::CastOp> {
       value =
           rewriter.createOrFold<plier::SignCastOp>(loc, signlessSrcType, value);
 
+    value = convertTensorElements(rewriter, loc, value, signlessDstType,
+                                  srcElem, dstElem);
+
     bool isSrcMemref = srcType.isa<mlir::MemRefType>();
     bool isDstMemref = dstType.isa<mlir::MemRefType>();