Optimize contigious layout arrays handling (#150)

Author: Hardcode84

numba_dpcomp/numba_dpcomp/mlir/tests/test_numpy.py

@@ -814,6 +814,26 @@ def py_func(a):
 
     assert_equal(py_func(arr), jit_func(arr))
 
+def test_contigious_layout_opt():
+    def py_func(a):
+        return a[0,1]
+
+    jit_func = njit(py_func)
+
+    a = np.array([[1,2],[3,4]])
+    b = a.T
+
+    with print_pass_ir([],['MakeStridedLayoutPass']):
+        assert_equal(py_func(a), jit_func(a))
+        ir = get_print_buffer()
+        assert ir.count('affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>') == 0, ir
+
+    with print_pass_ir([],['MakeStridedLayoutPass']):
+        assert_equal(py_func(b), jit_func(b))
+        ir = get_print_buffer()
+        assert ir.count('affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>') == 1, ir
+
+
 @parametrize_function_variants("a", [
     # 'np.array(1)', TODO zero rank arrays
     # 'np.array(2.5)',

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/lower_to_gpu.cpp

@@ -33,6 +33,7 @@
 #include <mlir/Dialect/Arithmetic/Transforms/Passes.h>
 #include <mlir/Dialect/GPU/ParallelLoopMapper.h>
 #include <mlir/Dialect/GPU/Passes.h>
+#include <mlir/Dialect/GPU/Utils.h>
 #include <mlir/Dialect/LLVMIR/LLVMDialect.h>
 #include <mlir/Dialect/Math/IR/Math.h>
 #include <mlir/Dialect/MemRef/IR/MemRef.h>
@@ -2593,6 +2594,13 @@ void MarkGpuArraysInputs::runOnOperation() {
   auto func = getOperation();
   auto funcType = func.getType();
 
+  mlir::OpBuilder builder(&getContext());
+  auto attrStr = builder.getStringAttr(kGpuArgAttr);
+  if (func->hasAttr(attrStr)) {
+    markAllAnalysesPreserved();
+    return;
+  }
+
   bool needAttr = false;
   llvm::SmallVector<bool> result;
   result.reserve(funcType.getNumInputs());
@@ -2606,10 +2614,8 @@ void MarkGpuArraysInputs::runOnOperation() {
   for (auto type : (func.getType().getInputs()))
     visitTypeRecursive(type, visitor);
 
-  if (needAttr) {
-    mlir::OpBuilder builder(&getContext());
-    func->setAttr(kGpuArgAttr, builder.getBoolArrayAttr(result));
-  }
+  if (needAttr)
+    func->setAttr(attrStr, builder.getBoolArrayAttr(result));
 
   markAllAnalysesPreserved();
 }
@@ -2921,6 +2927,31 @@ struct LowerGpuBuiltinsPass
     : public plier::RewriteWrapperPass<LowerGpuBuiltinsPass, void, void,
                                        LowerPlierCalls, LowerBuiltinCalls> {};
 
+class GpuLaunchSinkOpsPass
+    : public mlir::PassWrapper<GpuLaunchSinkOpsPass,
+                               mlir::OperationPass<void>> {
+public:
+  void runOnOperation() override {
+    using namespace mlir;
+
+    Operation *op = getOperation();
+    if (op->walk([](gpu::LaunchOp launch) {
+            auto isSinkingBeneficiary = [](mlir::Operation *op) -> bool {
+              return isa<arith::ConstantOp, ConstantOp, arith::SelectOp,
+                         arith::CmpIOp>(op);
+            };
+
+            // Pull in instructions that can be sunk
+            if (failed(
+                    sinkOperationsIntoLaunchOp(launch, isSinkingBeneficiary)))
+              return WalkResult::interrupt();
+
+            return WalkResult::advance();
+          }).wasInterrupted())
+      signalPassFailure();
+  }
+};
+
 static void commonOptPasses(mlir::OpPassManager &pm) {
   pm.addPass(mlir::createCanonicalizerPass());
   pm.addPass(mlir::createCSEPass());
@@ -2949,12 +2980,15 @@ static void populateLowerToGPUPipelineLow(mlir::OpPassManager &pm) {
   funcPM.addPass(mlir::createCanonicalizerPass());
   funcPM.addPass(std::make_unique<UnstrideMemrefsPass>());
   funcPM.addPass(mlir::createLowerAffinePass());
+
+  // TODO: mlir::gpu::GPUModuleOp pass
+  pm.addNestedPass<mlir::FuncOp>(mlir::arith::createArithmeticExpandOpsPass());
   commonOptPasses(funcPM);
   funcPM.addPass(std::make_unique<KernelMemrefOpsMovementPass>());
-
+  funcPM.addPass(std::make_unique<GpuLaunchSinkOpsPass>());
+  pm.addPass(mlir::createGpuKernelOutliningPass());
   pm.addPass(mlir::createSymbolDCEPass());
 
-  pm.addPass(mlir::createGpuKernelOutliningPass());
   pm.addNestedPass<mlir::FuncOp>(std::make_unique<GPULowerDefaultLocalSize>());
   pm.nest<mlir::gpu::GPUModuleOp>().addNestedPass<mlir::gpu::GPUFuncOp>(
       mlir::arith::createArithmeticExpandOpsPass());

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

@@ -1200,6 +1200,9 @@ struct ReshapeChangeLayout
   }
 };
 
+static constexpr llvm::StringLiteral
+    kContigiousArraysAttr("plier.contigious_arrays");
+
 struct MakeStridedLayoutPass
     : public mlir::PassWrapper<MakeStridedLayoutPass,
                                mlir::OperationPass<mlir::ModuleOp>> {
@@ -1210,23 +1213,48 @@ void MakeStridedLayoutPass::runOnOperation() {
   auto context = &getContext();
   auto mod = getOperation();
 
+  mlir::OpBuilder builder(mod);
+  auto loc = builder.getUnknownLoc();
+  auto attrStr = builder.getStringAttr(kContigiousArraysAttr);
+
+  llvm::SmallVector<bool> contigiousArrayArg;
+
+  auto isContigiousArrayArg = [&](unsigned i) {
+    if (contigiousArrayArg.empty())
+      return false;
+
+    assert(i < contigiousArrayArg.size());
+    return contigiousArrayArg[i];
+  };
+
   llvm::SmallVector<mlir::Type> newArgTypes;
   llvm::SmallVector<mlir::Type> newResTypes;
   llvm::SmallVector<mlir::Value> newOperands;
   for (auto func : mod.getOps<mlir::FuncOp>()) {
-    mlir::OpBuilder builder(func.body());
-    auto loc = builder.getUnknownLoc();
+    auto contAttr = func->getAttr(attrStr).dyn_cast_or_null<mlir::ArrayAttr>();
+    if (contAttr) {
+      auto contAttrRange = contAttr.getAsValueRange<mlir::BoolAttr>();
+      contigiousArrayArg.assign(contAttrRange.begin(), contAttrRange.end());
+    } else {
+      contigiousArrayArg.clear();
+    }
+
     auto funcType = func.getType();
     auto argTypes = funcType.getInputs();
     auto resTypes = funcType.getResults();
     newArgTypes.assign(argTypes.begin(), argTypes.end());
     newResTypes.assign(resTypes.begin(), resTypes.end());
-    bool hasBody = !func.getBody().empty();
+    auto &body = func.getBody();
+    bool hasBody = !body.empty();
+    if (hasBody)
+      builder.setInsertionPointToStart(&body.front());
+
     for (auto it : llvm::enumerate(argTypes)) {
       auto i = static_cast<unsigned>(it.index());
       auto type = it.value();
       auto memrefType = type.dyn_cast<mlir::MemRefType>();
-      if (!memrefType || !memrefType.getLayout().isIdentity())
+      if (!memrefType || isContigiousArrayArg(i) ||
+          !memrefType.getLayout().isIdentity())
         continue;
 
       auto rank = static_cast<unsigned>(memrefType.getRank());
@@ -1244,7 +1272,7 @@ void MakeStridedLayoutPass::runOnOperation() {
         newArgTypes[i] = newMemrefType;
 
         if (hasBody) {
-          auto arg = func.getBody().front().getArgument(i);
+          auto arg = body.front().getArgument(i);
           arg.setType(newMemrefType);
           auto dst =
               builder.create<plier::ChangeLayoutOp>(loc, memrefType, arg);
@@ -2116,6 +2144,63 @@ void PostPlierToLinalgPass::runOnOperation() {
   (void)mlir::applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
 }
 
+template <typename F>
+static void visitTypeRecursive(mlir::Type type, F &&visitor) {
+  if (auto tupleType = type.dyn_cast<mlir::TupleType>()) {
+    for (auto t : tupleType.getTypes())
+      visitTypeRecursive(t, std::forward<F>(visitor));
+  } else {
+    visitor(type);
+  }
+}
+
+static bool isContigiousArray(mlir::Type type) {
+  auto pyType = type.dyn_cast<plier::PyType>();
+  if (!pyType)
+    return false;
+
+  auto name = pyType.getName();
+  auto desc = parseArrayDesc(name);
+  if (!desc)
+    return false;
+
+  return desc->layout == ArrayLayout::C;
+}
+
+struct MarkContigiousArraysPass
+    : public mlir::PassWrapper<MarkContigiousArraysPass,
+                               mlir::OperationPass<mlir::FuncOp>> {
+  void runOnOperation() override {
+    auto func = getOperation();
+    auto funcType = func.getType();
+
+    mlir::OpBuilder builder(&getContext());
+    auto attrStr = builder.getStringAttr(kContigiousArraysAttr);
+    if (func->hasAttr(attrStr)) {
+      markAllAnalysesPreserved();
+      return;
+    }
+
+    bool needAttr = false;
+    llvm::SmallVector<bool> result;
+    result.reserve(funcType.getNumInputs());
+
+    auto visitor = [&](mlir::Type type) {
+      auto res = isContigiousArray(type);
+      result.emplace_back(res);
+      needAttr = needAttr || res;
+    };
+
+    for (auto type : (func.getType().getInputs()))
+      visitTypeRecursive(type, visitor);
+
+    if (needAttr)
+      func->setAttr(attrStr, builder.getBoolArrayAttr(result));
+
+    markAllAnalysesPreserved();
+  }
+};
+
 template <typename Op>
 struct ConvertAlloc : public mlir::OpConversionPattern<Op> {
   using mlir::OpConversionPattern<Op>::OpConversionPattern;
@@ -2601,6 +2686,7 @@ struct FixDeallocPlacementPass
                                        void, FixDeallocPlacement> {};
 
 static void populatePlierToLinalgGenPipeline(mlir::OpPassManager &pm) {
+  pm.addNestedPass<mlir::FuncOp>(std::make_unique<MarkContigiousArraysPass>());
   pm.addPass(std::make_unique<PlierToLinalgPass>());
   pm.addPass(mlir::createCanonicalizerPass());
   pm.addPass(std::make_unique<NumpyCallsLoweringPass>());