[MLIR] Numpy empty and sum axis

mlir-compiler/mlir-compiler/src/pipelines/plier_to_linalg.cpp

@@ -133,8 +133,12 @@ bool is_int(mlir::Type type)
     return type.isa<mlir::IntegerType>();
 }
 
-mlir::LogicalResult lower_prange(plier::PyCallOp op, llvm::ArrayRef<mlir::Value> operands, mlir::PatternRewriter& rewriter)
+mlir::LogicalResult lower_prange(plier::PyCallOp op, llvm::ArrayRef<mlir::Value> operands, llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>> kwargs, mlir::PatternRewriter& rewriter)
 {
+    if (!kwargs.empty())
+    {
+        return mlir::failure();
+    }
     if ((operands.size() < 1 || operands.size() > 3) ||
         !llvm::all_of(operands, [](mlir::Value val) { return is_int(val.getType());}))
     {
@@ -177,21 +181,21 @@ struct CallLowerer
 {
     mlir::LogicalResult operator()(
         plier::PyCallOp op, llvm::StringRef name, llvm::ArrayRef<mlir::Value> args,
-        mlir::PatternRewriter& rewriter)
+        llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>> kwargs, mlir::PatternRewriter& rewriter)
     {
-        using func_t = mlir::LogicalResult(*)(plier::PyCallOp, llvm::ArrayRef<mlir::Value>, mlir::PatternRewriter&);
+        using func_t = mlir::LogicalResult(*)(plier::PyCallOp, llvm::ArrayRef<mlir::Value>, llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>>, mlir::PatternRewriter&);
         std::pair<llvm::StringRef, func_t> handlers[] = {
             {"numba.prange", lower_prange},
         };
         for (auto& handler : handlers)
         {
             if (handler.first == name)
             {
-                return handler.second(op, args, rewriter);
+                return handler.second(op, args, kwargs, rewriter);
             }
         }
 
-        if (auto result = linalg_resolver.rewrite(name, op.getLoc(), rewriter, args))
+        if (auto result = linalg_resolver.rewrite(name, op.getLoc(), rewriter, args, kwargs))
         {
             assert(result->size() == op->getNumResults());
             rerun_std_pipeline(op);
@@ -206,7 +210,7 @@ struct CallLowerer
             return mlir::success();
         }
 
-        if (name == "len" && check_numpy_args(args, 1))
+        if (name == "len" && check_numpy_args(args, 1) && kwargs.empty())
         {
             auto loc = op.getLoc();
             mlir::Value dim = rewriter.create<mlir::DimOp>(loc, args[0], 0);
@@ -219,7 +223,6 @@ struct CallLowerer
     }
 
 private:
-
     PyLinalgResolver linalg_resolver;
 };
 
@@ -436,7 +439,7 @@ struct SetitemOpLowering : public mlir::OpRewritePattern<T>
             mlir::OpBuilder::InsertionGuard g(rewriter);
             if (auto parent_op = target.getDefiningOp())
             {
-                rewriter.setInsertionPoint(parent_op);
+                rewriter.setInsertionPointAfter(parent_op);
             }
             else
             {
@@ -456,6 +459,7 @@ struct SetitemOpLowering : public mlir::OpRewritePattern<T>
                     }
                     else
                     {
+                        mlir::OpBuilder::InsertionGuard g(rewriter);
                         rewriter.setInsertionPoint(use_op);
                         auto new_val = rewriter.create<mlir::TensorLoadOp>(use_op->getLoc(), memref);
                         rewriter.updateRootInPlace(use_op, [&]()
@@ -602,6 +606,7 @@ struct LowerLinalgPass :
         mlir::DialectRegistry &registry) const override
     {
         registry.insert<mlir::StandardOpsDialect>();
+        registry.insert<mlir::tensor::TensorDialect>();
         registry.insert<mlir::linalg::LinalgDialect>();
         registry.insert<mlir::scf::SCFDialect>();
         registry.insert<mlir::AffineDialect>();
@@ -686,6 +691,7 @@ void PostLinalgOptPass::runOnOperation()
 void populate_plier_to_linalg_gen_pipeline(mlir::OpPassManager& pm)
 {
     pm.addPass(std::make_unique<PlierToLinalgPass>());
+    pm.addPass(mlir::createSymbolDCEPass());
 }
 
 void populate_plier_to_linalg_opt_pipeline(mlir::OpPassManager& pm)
@@ -708,6 +714,7 @@ void populate_plier_to_linalg_opt_pipeline(mlir::OpPassManager& pm)
 
     pm.addPass(std::make_unique<LowerLinalgPass>());
     pm.addPass(std::make_unique<PostLinalgOptPass>());
+    pm.addPass(mlir::createSymbolDCEPass());
 }
 }
 

mlir-compiler/mlir-compiler/src/pipelines/plier_to_std.cpp

@@ -1156,8 +1156,12 @@ struct FoldTupleGetitem : public mlir::OpRewritePattern<Op>
     }
 };
 
-mlir::LogicalResult lower_range(plier::PyCallOp op, llvm::ArrayRef<mlir::Value> operands, mlir::PatternRewriter& rewriter)
+mlir::LogicalResult lower_range(plier::PyCallOp op, llvm::ArrayRef<mlir::Value> operands, llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>> kwargs, mlir::PatternRewriter& rewriter)
 {
+    if (!kwargs.empty())
+    {
+        return mlir::failure();
+    }
     if ((operands.size() < 1 || operands.size() > 3) ||
         !llvm::all_of(operands, [](mlir::Value val) { return is_int(val.getType());}))
     {
@@ -1191,8 +1195,12 @@ mlir::LogicalResult lower_range(plier::PyCallOp op, llvm::ArrayRef<mlir::Value>
     return mlir::success();
 }
 
-mlir::LogicalResult lower_len(plier::PyCallOp op, llvm::ArrayRef<mlir::Value> operands, mlir::PatternRewriter& rewriter)
+mlir::LogicalResult lower_len(plier::PyCallOp op, llvm::ArrayRef<mlir::Value> operands, llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>> kwargs, mlir::PatternRewriter& rewriter)
 {
+    if (!kwargs.empty())
+    {
+        return mlir::failure();
+    }
     if (operands.size() != 1)
     {
         return mlir::failure();
@@ -1210,8 +1218,12 @@ mlir::LogicalResult lower_len(plier::PyCallOp op, llvm::ArrayRef<mlir::Value> op
     return mlir::success();
 }
 
-mlir::LogicalResult lower_bool_cast(plier::PyCallOp op, llvm::ArrayRef<mlir::Value> operands, mlir::PatternRewriter& rewriter)
+mlir::LogicalResult lower_bool_cast(plier::PyCallOp op, llvm::ArrayRef<mlir::Value> operands, llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>> kwargs, mlir::PatternRewriter& rewriter)
 {
+    if (!kwargs.empty())
+    {
+        return mlir::failure();
+    }
     if (operands.size() != 1)
     {
         return mlir::failure();
@@ -1250,8 +1262,12 @@ mlir::FuncOp get_lib_symbol(
 
 mlir::LogicalResult lower_math_func(
     plier::PyCallOp op, llvm::StringRef name, llvm::ArrayRef<mlir::Value> args,
-    mlir::PatternRewriter& rewriter)
+    llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>> kwargs, mlir::PatternRewriter& rewriter)
 {
+    if (!kwargs.empty())
+    {
+        return mlir::failure();
+    }
     auto ret_type = map_plier_type(op.getType());
     auto valid_type = [&](mlir::Type type)
     {
@@ -1285,14 +1301,14 @@ mlir::LogicalResult lower_math_func(
 struct CallLowerer
 {
     mlir::LogicalResult operator()(plier::PyCallOp op, llvm::StringRef name,
-        llvm::ArrayRef<mlir::Value> args, mlir::PatternRewriter& rewriter)
+        llvm::ArrayRef<mlir::Value> args, llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>> kwargs, mlir::PatternRewriter& rewriter)
     {
-        if (mlir::succeeded(lower_math_func(op, name, args, rewriter)))
+        if (mlir::succeeded(lower_math_func(op, name, args, kwargs, rewriter)))
         {
             return mlir::success();
         }
 
-        using func_t = mlir::LogicalResult(*)(plier::PyCallOp, llvm::ArrayRef<mlir::Value>, mlir::PatternRewriter&);
+        using func_t = mlir::LogicalResult(*)(plier::PyCallOp, llvm::ArrayRef<mlir::Value>, llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>>, mlir::PatternRewriter&);
         std::pair<llvm::StringRef, func_t> handlers[] = {
             {"bool", lower_bool_cast},
             {"range", lower_range},
@@ -1302,7 +1318,7 @@ struct CallLowerer
         {
             if (handler.first == name)
             {
-                return handler.second(op, args, rewriter);
+                return handler.second(op, args, kwargs, rewriter);
             }
         }