speed up developing ofstream in mlir codegen (#10152)

Signed-off-by: daquexian <daquexian566@gmail.com> Co-authored-by: oneflow-ci-bot <ci-bot@oneflow.org> Co-authored-by: daquexian <daquexian566@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Oneflow-Inc · Apr 17, 2023 · aa22da4 · aa22da4
1 parent d78d84d
commit aa22da4
Show file tree

Hide file tree

Showing 14 changed files with 154 additions and 7 deletions.
diff --git a/oneflow/api/python/framework/tensor.cpp b/oneflow/api/python/framework/tensor.cpp
@@ -400,8 +400,10 @@ static PyObject* PyTensorObject_check_meta_consistency(PyObject* self, PyObject*
 static PyObject* PyTensorObject_data_ptr(PyObject* self, PyObject* unused) {
   HANDLE_ERRORS
   const auto& t = PyTensor_Unpack(self);
+  const std::shared_ptr<LocalTensor> local_tensor =
+      t->is_local() ? ASSERT_PTR(t->AsLocalTensor()) : ASSERT_PTR(t->cur_rank_phy_tensor());
   return functional::CastToPyObject(
-      reinterpret_cast<int64_t>(ASSERT(GetTensorDataPtr(ASSERT_PTR(t->AsLocalTensor())))));
+      reinterpret_cast<int64_t>(ASSERT(GetTensorDataPtr(local_tensor))));
   END_HANDLE_ERRORS
 }
 

diff --git a/oneflow/ir/include/OneFlow/OneFlowPDLLPatterns.h b/oneflow/ir/include/OneFlow/OneFlowPDLLPatterns.h
@@ -21,6 +21,7 @@ namespace mlir {
 
 namespace oneflow {
 
+void populateAllocEliminationPatterns(RewritePatternSet& patterns);
 void populateForwardOpPatterns(RewritePatternSet& patterns);
 void populateNormalizationOpPatterns(RewritePatternSet& patterns);
 void populateFuseConv2DBatchNormPattern(RewritePatternSet& patterns);

diff --git a/oneflow/ir/include/OneFlow/OneFlowPasses.td b/oneflow/ir/include/OneFlow/OneFlowPasses.td
@@ -69,10 +69,15 @@ def WrapOpsToKernelLaunchPass : Pass<"wrap-ops-to-kernel-launch", "ModuleOp"> {
   let constructor = "mlir::oneflow::createWrapOpsToKernelLaunchPass()";
 }
 
+def EliminateAllocOpsPass : Pass<"eliminate-alloc-ops", "ModuleOp"> {
+  let summary = "";
+  let constructor = "mlir::oneflow::createEliminateAllocOpsPass()";
+  let dependentDialects = ["pdl_interp::PDLInterpDialect", "pdl::PDLDialect"];
+}
+
 def AppendOneFlowStreamPass : Pass<"append-ofstream", "ModuleOp"> {
   let summary = "append oneflow stream to gpu function arguments";
   let constructor = "mlir::oneflow::createAppendOneFlowStreamPass()";
-  let dependentDialects = ["pdl_interp::PDLInterpDialect", "pdl::PDLDialect"];
 }
 
 def FuseIntoExistingOpPass : Pass<"fuse-into-existing-op", "ModuleOp"> {

diff --git a/oneflow/ir/include/OneFlow/Passes.h b/oneflow/ir/include/OneFlow/Passes.h
@@ -34,6 +34,7 @@ limitations under the License.
 #include "OneFlow/Transform/AggregateOps.h"
 #include "OneFlow/Transform/FuncOps.h"
 #include "OneFlow/Transform/CSEWithAttributesIgnored.h"
+#include "OneFlow/Transform/EliminateAllocOps.h"
 #include "OneFlow/Transform/AppendOneFlowStream.h"
 
 #ifdef WITH_MLIR_CUDA_CODEGEN

diff --git a/oneflow/ir/include/OneFlow/Transform/EliminateAllocOps.h b/oneflow/ir/include/OneFlow/Transform/EliminateAllocOps.h
@@ -0,0 +1,31 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#ifndef ONEFLOW_IR_INCLUDE_ONEFLOW_TRANSFORM_ELIMINATE_ALLOC_OPS_H_
+#define ONEFLOW_IR_INCLUDE_ONEFLOW_TRANSFORM_ELIMINATE_ALLOC_OPS_H_
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+namespace oneflow {
+
+std::unique_ptr<mlir::Pass> createEliminateAllocOpsPass();
+
+}  // namespace oneflow
+
+}  // namespace mlir
+
+#endif  // ONEFLOW_IR_INCLUDE_ONEFLOW_TRANSFORM_ELIMINATE_ALLOC_OPS_H_
diff --git a/oneflow/ir/lib/OneFlow/CMakeLists.txt b/oneflow/ir/lib/OneFlow/CMakeLists.txt
@@ -60,6 +60,7 @@ oneflow_add_mlir_dialect_library(
   Transform/AutoNhwc.cpp
   Transform/ConvertInferenceOp.cpp
   Transform/AggregateOps.cpp
+  Transform/EliminateAllocOps.cpp
   Transform/FuncOps.cpp
   Transform/CSEWithAttributesIgnored.cpp
   Transform/GroupMatMulOps.cpp

diff --git a/oneflow/ir/lib/OneFlow/Conversion/OneFlowToTosa.cpp b/oneflow/ir/lib/OneFlow/Conversion/OneFlowToTosa.cpp
@@ -623,14 +623,10 @@ struct CastInputConversion final : public OpRewritePattern<InputOp> {
         if (isSignLessTensorOrOther(cast.getResult(0).getType())) { return failure(); }
       }
     }
-    LOG(ERROR) << "ok4";
     InputOp cloned = rewriter.create<InputOp>(op->getLoc(), op.getResultTypes(), op->getOperands(),
                                               op->getAttrs());
-    auto m = op->getParentOp();
-    m->dump();
     rewriter.replaceOpWithNewOp<UnrealizedConversionCastOp>(
         op, convertToSignless(getContext(), op.getOutput().getType()), cloned.getOutput());
-    m->dump();
     return success();
   }
 };

diff --git a/oneflow/ir/lib/OneFlow/PDLL/AllocEliminationPatterns.cpp b/oneflow/ir/lib/OneFlow/PDLL/AllocEliminationPatterns.cpp
@@ -0,0 +1,41 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/PDL/IR/PDL.h"
+#include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "OneFlow/OneFlowPDLLPatterns.h"
+#include "mlir/IR/Value.h"
+
+using namespace mlir;
+
+#include "oneflow/ir/lib/OneFlow/PDLL/AllocEliminationPatterns.h.inc"
+
+namespace mlir {
+
+namespace oneflow {
+
+void populateAllocEliminationPatterns(RewritePatternSet& patterns) {
+  populateGeneratedPDLLPatterns(patterns);
+}
+
+}  // namespace oneflow
+
+}  // namespace mlir
diff --git a/oneflow/ir/lib/OneFlow/PDLL/AllocEliminationPatterns.pdll b/oneflow/ir/lib/OneFlow/PDLL/AllocEliminationPatterns.pdll
@@ -0,0 +1,15 @@
+#include "OneFlow/OneFlowOps.td"
+
+Constraint IsFuncArguments(value: Value) [{
+  return success(llvm::dyn_cast<mlir::BlockArgument>(value));
+}];
+
+Pattern {
+  let alloc = op<memref.alloc>();
+  let copy = op<memref.copy>(alloc.0, arg: IsFuncArguments);
+
+  rewrite alloc with {
+    erase copy;
+    replace alloc with arg;
+  };
+}
diff --git a/oneflow/ir/lib/OneFlow/PDLL/CMakeLists.txt b/oneflow/ir/lib/OneFlow/PDLL/CMakeLists.txt
@@ -1,3 +1,5 @@
+add_mlir_pdll_library(MLIROneFlowPDLLAllocEliminaionPatternsIncGen AllocEliminationPatterns.pdll
+                      AllocEliminationPatterns.h.inc)
 add_mlir_pdll_library(MLIROneFlowPDLLForwardOpPatternsIncGen ForwardOpPatterns.pdll
                       ForwardOpPatterns.h.inc)
 add_mlir_pdll_library(MLIROneFlowPDLLNormalizationPatternsIncGen NormalizationPatterns.pdll
@@ -8,11 +10,13 @@ add_mlir_pdll_library(MLIROneFlowPDLLFuseOpsWithBackwardImplPatternsIncGen
                       FuseOpsWithBackwardImplPattern.pdll FuseOpsWithBackwardImplPattern.h.inc)
 oneflow_add_mlir_dialect_library(
   MLIROneFlowPDLLPatterns
+  AllocEliminationPatterns.cpp
   ForwardOpPatterns.cpp
   NormalizationPatterns.cpp
   FuseConv2DBatchNormPattern.cpp
   FuseOpsWithBackwardImplPattern.cpp
   DEPENDS
+  MLIROneFlowPDLLAllocEliminaionPatternsIncGen
   MLIROneFlowPDLLForwardOpPatternsIncGen
   MLIROneFlowPDLLNormalizationPatternsIncGen
   MLIROneFlowPDLLFuseConv2DBatchNormPatternIncGen

diff --git a/oneflow/ir/lib/OneFlow/Passes.cpp b/oneflow/ir/lib/OneFlow/Passes.cpp
@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "OneFlow/Transform/EliminateAllocOps.h"
 #include "OneFlow/Transform/AppendOneFlowStream.h"
 #include "oneflow/ir/oneflow-translate/include/OneFlow/MLIROneFlowTranslation.h"
 #include "oneflow/core/kernel/cuda_graph_support.h"
@@ -1029,6 +1030,7 @@ void AddLowerToLinalgMemRefPasses(PassManager& pm) {
   pm.addNestedPass<func::FuncOp>(createTensorBufferizePass());      // tensor-bufferize
   pm.addPass(func::createFuncBufferizePass());                      // func-bufferize
   pm.addPass(bufferization::createBufferResultsToOutParamsPass());  // buffer-results-to-out-params
+  pm.addPass(mlir::oneflow::createEliminateAllocOpsPass());         // eliminate-alloc-ops
   pm.addPass(createCanonicalizerPass());                            // canonicalize
   pm.addNestedPass<func::FuncOp>(
       mlir::bufferization::createFinalizingBufferizePass());  // finalizing-bufferize

diff --git a/oneflow/ir/lib/OneFlow/Transform/EliminateAllocOps.cpp b/oneflow/ir/lib/OneFlow/Transform/EliminateAllocOps.cpp
@@ -0,0 +1,40 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "OneFlow/OneFlowPDLLPatterns.h"
+#include "OneFlow/Passes.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir {
+namespace oneflow {
+
+namespace {
+class EliminateAllocOpsPass : public EliminateAllocOpsPassBase<EliminateAllocOpsPass> {
+  void runOnOperation() override {
+    Operation* op = getOperation();
+    RewritePatternSet patterns(op->getContext());
+    mlir::oneflow::populateAllocEliminationPatterns(patterns);
+    (void)applyPatternsAndFoldGreedily(op, std::move(patterns));
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<Pass> createEliminateAllocOpsPass() {
+  return std::make_unique<EliminateAllocOpsPass>();
+}
+
+}  // namespace oneflow
+}  // namespace mlir
diff --git a/oneflow/ir/oneflow-opt/oneflow-opt.cpp b/oneflow/ir/oneflow-opt/oneflow-opt.cpp
@@ -78,6 +78,7 @@ int32_t main(int32_t argc, char** argv) {
   mlir::registerOutlineJitFunctionPassPass();
   mlir::oneflow::registerCSEPasses(global_cse_state);
   mlir::registerFuseForwardOpsPass();
+  mlir::registerEliminateAllocOpsPassPass();
   mlir::registerFuseIntoExistingOpPassPass();
   mlir::registerFuseNormalizationOpsPass();
   mlir::registerFuseOpsWithBackwardImplPass();

diff --git a/python/oneflow/test/tensor/test_data_ptr.py b/python/oneflow/test/tensor/test_data_ptr.py
@@ -18,8 +18,8 @@
 import oneflow.unittest
 
 
-@flow.unittest.skip_unless_1n1d()
 class TestDataPtr(unittest.TestCase):
+    @flow.unittest.skip_unless_1n1d()
     def test_equality(test_case):
         x = flow.ones(2, 3)
         y = flow.ones(2, 3)
@@ -31,6 +31,13 @@ def test_equality(test_case):
         x[:] = 2
         test_case.assertEqual(x_ptr, x.data_ptr())
 
+    @flow.unittest.skip_unless_1n2d()
+    def test_global_tensor(test_case):
+        x = flow.randn(
+            2, 3, placement=flow.placement.all("cpu"), sbp=flow.sbp.broadcast
+        )
+        test_case.assertEqual(x.data_ptr(), x.to_local().data_ptr())
+
 
 if __name__ == "__main__":
     unittest.main()