Support reshape with negative val (#188)

Author: Hardcode84

azure-pipelines.yml

@@ -41,7 +41,7 @@ jobs:
   - script: |
       call "C:\Miniconda\Scripts\activate"
       call cd numba_dpcomp
-      pytest -n1 -vv --capture=tee-sys -rX
+      pytest -n1 -vv --capture=tee-sys -rXF
     displayName: 'Tests'
 
   - script: |
@@ -126,7 +126,7 @@ jobs:
       export OCL_ICD_FILENAMES=libintelocl.so
       export SYCL_DEVICE_FILTER=opencl:cpu
       export NUMBA_DISABLE_PERFORMANCE_WARNINGS=1
-      pytest -n1 -vv --capture=tee-sys -rX
+      pytest -n1 -vv --capture=tee-sys -rXF
     displayName: 'Tests'
 
   - script: |
@@ -196,7 +196,7 @@ jobs:
       source /usr/local/miniconda/bin/activate
       cd numba_dpcomp
       conda activate test_env
-      pytest -n1 -vv --capture=tee-sys -rX
+      pytest -n1 -vv --capture=tee-sys -rXF
     displayName: 'Tests'
 
   - script: |

numba_dpcomp/numba_dpcomp/mlir/linalg_builder.py

@@ -65,6 +65,9 @@ def __rmul__(self, o):
     def __truediv__(self, o):
         return self._binop(self._context, self._ssa_val, o, "/")
 
+    def __floordiv__(self, o):
+        return self._binop(self._context, self._ssa_val, o, "//")
+
     def __lt__(self, o):
         return self._binop(self._context, self._ssa_val, o, "lt")
 

numba_dpcomp/numba_dpcomp/mlir/numpy/funcs.py

@@ -587,11 +587,22 @@ def reshape_impl(builder, arg, *new_shape):
     if len(new_shape) == 1:
         new_shape = new_shape[0]
 
-    # TODO: better check
     if isinstance(new_shape, tuple):
-        for s in new_shape:
+        neg_index = None
+        for i, s in enumerate(new_shape):
             if isinstance(s, int) and s < 0:
-                return  # not supported for now
+                assert neg_index is None
+                neg_index = i
+        if neg_index is not None:
+            size = 1
+            for i, s in enumerate(new_shape):
+                if i != neg_index:
+                    size = size * s
+
+            size = size_impl(builder, arg) // size
+            new_shape = list(new_shape)
+            new_shape[neg_index] = size
+            new_shape = tuple(new_shape)
 
     return builder.reshape(arg, new_shape)
 

numba_dpcomp/numba_dpcomp/mlir/tests/test_numba_parfor.py

@@ -95,7 +95,6 @@ def _gen_tests():
         "test_simple18",  # np.linalg.svd
         "test_linspace",  # np.linspace
         "test_std",  # array.std
-        "test_reshape_with_neg_one",  # unsupported reshape
         "test_mvdot",  # np.dot unsupported args
         "test_array_tuple_concat",  # tuple concat
         "test_namedtuple1",  # namedtuple support
@@ -141,7 +140,6 @@ def _gen_tests():
         "test_issue6102",  # list support
         "test_oversized_tuple_as_arg_to_kernel",  # UnsupportedParforsError not raised
         "test_issue5942_2",  # invalid result
-        "test_reshape_with_large_neg",  # unsupported reshape
         "test_parfor_ufunc_typing",  # np.isinf
         "test_issue_5098",  # list support and more
         "test_parfor_slice27",  # Literal return issue

numba_dpcomp/numba_dpcomp/mlir/tests/test_numpy.py

@@ -1006,7 +1006,13 @@ def py_func(a):
         "lambda a: a.reshape((1, a.size))",
         "lambda a: a.reshape(1, a.size)",
         "lambda a: a.reshape((1, a.size, 1))",
+        "lambda a: a.reshape((-1, a.size, 1))",
+        "lambda a: a.reshape((1, -1, 1))",
+        "lambda a: a.reshape((1, a.size, -1))",
         "lambda a: a.reshape(1, a.size, 1)",
+        "lambda a: a.reshape(-1, a.size, 1)",
+        "lambda a: a.reshape(1, -1, 1)",
+        "lambda a: a.reshape(1, a.size, -1)",
         "lambda a: np.reshape(a, a.size)",
         "lambda a: np.reshape(a, (a.size,))",
         "lambda a: np.reshape(a, (a.size,1))",

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/py_linalg_resolver.cpp

@@ -232,20 +232,6 @@ static llvm::Optional<py::object> getPyLiteral(mlir::Attribute attr) {
   return {};
 }
 
-static llvm::Optional<py::object> makePyLiteral(mlir::Value val) {
-  assert(val);
-  if (auto literal = val.getType().dyn_cast<plier::LiteralType>())
-    return getPyLiteral(literal.getValue());
-
-  if (auto cast = val.getDefiningOp<plier::SignCastOp>())
-    val = cast.value();
-
-  if (auto attr = plier::getConstVal<mlir::Attribute>(val))
-    return getPyLiteral(attr);
-
-  return {};
-}
-
 static mlir::Value doCast(mlir::OpBuilder &builder, mlir::Location loc,
                           mlir::Value val, mlir::Type type) {
   if (val.getType() != type)
@@ -284,7 +270,7 @@ struct PyLinalgResolver::Context {
     if (auto typevar = type.dyn_cast<plier::TypeVar>())
       return createType(typevar.getType());
 
-    if (auto literal = makePyLiteral(value))
+    if (auto literal = makePyLiteral(context, value))
       return *literal;
 
     auto ret = var(context, wrapMlir(value));
@@ -400,6 +386,32 @@ struct PyLinalgResolver::Context {
 
     return doCast(builder, loc, unwrapVal(loc, builder, obj), resultType);
   }
+
+private:
+  llvm::Optional<py::object> makePyLiteral(py::capsule context,
+                                           mlir::Value val) {
+    assert(val);
+    if (auto literal = val.getType().dyn_cast<plier::LiteralType>())
+      return getPyLiteral(literal.getValue());
+
+    if (auto buildTuple = val.getDefiningOp<plier::BuildTupleOp>()) {
+      auto args = buildTuple.args();
+      auto count = static_cast<unsigned>(args.size());
+      py::tuple ret(count);
+      for (auto i : llvm::seq(0u, count))
+        ret[i] = createVar(context, args[i]);
+
+      return ret;
+    }
+
+    if (auto cast = val.getDefiningOp<plier::SignCastOp>())
+      val = cast.value();
+
+    if (auto attr = plier::getConstVal<mlir::Attribute>(val))
+      return getPyLiteral(attr);
+
+    return {};
+  }
 };
 
 namespace {
@@ -1653,13 +1665,19 @@ static py::object getitemImpl(py::capsule context, py::capsule ssaVal,
 template <typename Op>
 static mlir::Value binopFunc(mlir::Location loc, mlir::OpBuilder &builder,
                              mlir::Value lhs, mlir::Value rhs) {
-  return builder.create<Op>(loc, lhs, rhs);
+  auto lhsVar = doSignCast(builder, loc, lhs);
+  auto rhsVar = doSignCast(builder, loc, rhs);
+  auto res = builder.create<Op>(loc, lhsVar, rhsVar);
+  return doSignCast(builder, loc, res, lhs.getType());
 }
 
 template <typename Op>
 static mlir::Value rbinopFunc(mlir::Location loc, mlir::OpBuilder &builder,
                               mlir::Value lhs, mlir::Value rhs) {
-  return builder.create<Op>(loc, rhs, lhs);
+  auto lhsVar = doSignCast(builder, loc, lhs);
+  auto rhsVar = doSignCast(builder, loc, rhs);
+  auto res = builder.create<Op>(loc, rhsVar, lhsVar);
+  return doSignCast(builder, loc, res, lhs.getType());
 }
 
 static mlir::Value binopFuncIdiv(mlir::Location loc, mlir::OpBuilder &builder,
@@ -1669,6 +1687,14 @@ static mlir::Value binopFuncIdiv(mlir::Location loc, mlir::OpBuilder &builder,
   return builder.create<mlir::arith::DivFOp>(loc, lhsVar, rhsVar);
 }
 
+static mlir::Value binopFFloorDiv(mlir::Location loc, mlir::OpBuilder &builder,
+                                  mlir::Value lhs, mlir::Value rhs) {
+  auto lhsVar = doCast(builder, loc, lhs, builder.getF64Type());
+  auto rhsVar = doCast(builder, loc, rhs, builder.getF64Type());
+  auto res = builder.create<mlir::arith::DivFOp>(loc, lhsVar, rhsVar);
+  return builder.create<mlir::math::FloorOp>(loc, res);
+}
+
 template <mlir::arith::CmpIPredicate Pred>
 static mlir::Value binopCmpI(mlir::Location loc, mlir::OpBuilder &builder,
                              mlir::Value lhs, mlir::Value rhs) {
@@ -1713,6 +1739,7 @@ static py::object binopImpl(py::capsule context, py::capsule ssaVal,
        &rbinopFunc<mlir::arith::SubFOp>},
       {"*", &binopFunc<mlir::arith::MulIOp>, &binopFunc<mlir::arith::MulFOp>},
       {"/", &binopFuncIdiv, &binopFunc<mlir::arith::DivFOp>},
+      {"//", &binopFunc<mlir::arith::DivSIOp>, &binopFFloorDiv},
 
       {"lt", &binopCmpI<mlir::arith::CmpIPredicate::slt>,
        &binopCmpF<mlir::arith::CmpFPredicate::OLT>},