[python] Fix literal return issue (#192)

Author: Hardcode84

numba_dpcomp/numba_dpcomp/mlir/tests/test_basic.py

@@ -67,6 +67,15 @@ def py_func(a):
     assert_equal(py_func(val), jit_func(val))
 
 
+@pytest.mark.parametrize("val", _test_values)
+def test_ret_literal(val):
+    def py_func():
+        return val
+
+    jit_func = njit(py_func)
+    assert_equal(py_func(), jit_func())
+
+
 @parametrize_function_variants(
     "py_func",
     [

numba_dpcomp/numba_dpcomp/mlir/tests/test_numba_parfor.py

@@ -75,7 +75,6 @@ def _gen_tests():
         "test_unsigned_refusal_to_vectorize",  # Need to hook asm checks
         "test_vectorizer_fastmath_asm",  # Need to hook asm checks
         "test_kde_example",  # List suport
-        "test_prange27",  # Literal return issue
         "test_simple01",  # Empty shape not failed
         "test_kmeans",  # List suport
         "test_ndarray_fill",  # array.fill
@@ -140,7 +139,6 @@ def _gen_tests():
         "test_issue5942_2",  # invalid result
         "test_parfor_ufunc_typing",  # np.isinf
         "test_issue_5098",  # list support and more
-        "test_parfor_slice27",  # Literal return issue
         "test_ufunc_expr",  # np.bitwise_and(
         "test_parfor_generate_fuse",  # operand #0 does not dominate this use
         "test_parfor_slice7",  # array.transpose

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/lowering.cpp

@@ -272,17 +272,21 @@ struct PlierLowerer final {
       return builder.create<plier::ArgOp>(
           getCurrentLoc(), index, target.attr("name").cast<std::string>());
     }
-    if (py::isinstance(value, insts.Expr)) {
+
+    if (py::isinstance(value, insts.Expr))
       return lowerExpr(value);
-    }
-    if (py::isinstance(value, insts.Var)) {
+
+    if (py::isinstance(value, insts.Var))
       return loadvar(value);
-    }
-    if (py::isinstance(value, insts.Const)) {
+
+    if (py::isinstance(value, insts.Const))
       return getConst(value.attr("value"));
-    }
+
     if (py::isinstance(value, insts.Global) ||
         py::isinstance(value, insts.FreeVar)) {
+      auto constVal = getConstOrNull(value.attr("value"));
+      if (constVal)
+        return constVal;
       auto name = value.attr("name").cast<std::string>();
       return builder.create<plier::GlobalOp>(getCurrentLoc(), name);
     }
@@ -542,7 +546,7 @@ struct PlierLowerer final {
     builder.create<mlir::cf::BranchOp>(getCurrentLoc(), mlir::None, block);
   }
 
-  mlir::Value getConst(py::handle val) {
+  mlir::Value getConstOrNull(py::handle val) {
     auto getVal = [&](mlir::Attribute attr) {
       return builder.create<plier::ConstOp>(getCurrentLoc(), attr);
     };
@@ -558,8 +562,15 @@ struct PlierLowerer final {
     if (py::isinstance<py::none>(val))
       return getVal(builder.getUnitAttr());
 
-    plier::reportError(llvm::Twine("get_const unhandled type \"") +
-                       py::str(val.get_type()).cast<std::string>() + "\"");
+    return {};
+  }
+
+  mlir::Value getConst(py::handle val) {
+    auto ret = getConstOrNull(val);
+    if (!ret)
+      plier::reportError(llvm::Twine("get_const unhandled type \"") +
+                         py::str(val.get_type()).cast<std::string>() + "\"");
+    return ret;
   }
 
   mlir::FunctionType getFuncType(py::handle fnargs, py::handle restype) {

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/plier_to_std.cpp

@@ -1385,13 +1385,59 @@ void PlierToStdPass::runOnOperation() {
     signalPassFailure();
 }
 
+struct ConvertLiteralTypesPass
+    : public mlir::PassWrapper<PlierToStdPass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+  virtual void
+  getDependentDialects(mlir::DialectRegistry &registry) const override {
+    registry.insert<mlir::func::FuncDialect>();
+  }
+
+  void runOnOperation() override {
+    mlir::TypeConverter typeConverter;
+    // Convert unknown types to itself
+    typeConverter.addConversion([](mlir::Type type) { return type; });
+
+    auto context = &getContext();
+    typeConverter.addConversion(
+        [](plier::LiteralType type) { return type.getValue().getType(); });
+
+    auto materializeCast =
+        [](mlir::OpBuilder &builder, mlir::Type type, mlir::ValueRange inputs,
+           mlir::Location loc) -> llvm::Optional<mlir::Value> {
+      if (inputs.size() == 1)
+        return builder
+            .create<mlir::UnrealizedConversionCastOp>(loc, type, inputs.front())
+            .getResult(0);
+
+      return llvm::None;
+    };
+    typeConverter.addArgumentMaterialization(materializeCast);
+    typeConverter.addSourceMaterialization(materializeCast);
+    typeConverter.addTargetMaterialization(materializeCast);
+
+    mlir::RewritePatternSet patterns(context);
+    mlir::ConversionTarget target(*context);
+
+    plier::populateControlFlowTypeConversionRewritesAndTarget(typeConverter,
+                                                              patterns, target);
+    plier::populateTupleTypeConversionRewritesAndTarget(typeConverter, patterns,
+                                                        target);
+
+    if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,
+                                                  std::move(patterns))))
+      signalPassFailure();
+  }
+};
+
 static void populatePlierToStdPipeline(mlir::OpPassManager &pm) {
   pm.addPass(mlir::createCanonicalizerPass());
   pm.addPass(std::make_unique<PlierToStdPass>());
   pm.addPass(mlir::createCanonicalizerPass());
   pm.addPass(std::make_unique<BuiltinCallsLoweringPass>());
   pm.addPass(plier::createForceInlinePass());
   pm.addPass(mlir::createSymbolDCEPass());
+  pm.addPass(std::make_unique<ConvertLiteralTypesPass>());
   pm.addPass(mlir::createCanonicalizerPass());
 }
 } // namespace