1 /* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #include <memory>
17 
18 #include "mlir-hlo/Dialect/gml_st/IR/gml_st_ops.h"
19 #include "mlir-hlo/Dialect/gml_st/transforms/pass_detail.h"
20 #include "mlir-hlo/Dialect/gml_st/transforms/passes.h"
21 #include "mlir/Dialect/Func/IR/FuncOps.h"
22 #include "mlir/Dialect/Linalg/IR/Linalg.h"
23 #include "mlir/Dialect/Linalg/Transforms/CodegenStrategy.h"
24 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
25 #include "mlir/Dialect/Vector/IR/VectorOps.h"
26 #include "mlir/Pass/PassManager.h"
27 
28 namespace mlir {
29 namespace gml_st {
30 
31 struct VectorizeGmlStLoopsPass
32     : public VectorizeGmlStLoopsPassBase<VectorizeGmlStLoopsPass> {
runOnOperationmlir::gml_st::VectorizeGmlStLoopsPass33   void runOnOperation() override {
34     auto funcOp = getOperation();
35     // Vectorize linalg.generic operations inside gml_st.for and gml_st.parallel
36     // loops.
37     OpPassManager dynamicPM("func.func");
38     linalg::CodegenStrategy strategy;
39     strategy.vectorize(linalg::GenericOp::getOperationName(),
40                        [](mlir::Operation *op) {
41                          auto generic = mlir::dyn_cast<linalg::GenericOp>(op);
42                          if (!generic) return failure();
43                          if (op->getParentOfType<ForOp>() ||
44                              op->getParentOfType<ParallelOp>()) {
45                            return success();
46                          }
47                          return failure();
48                        });
49     strategy.configurePassPipeline(dynamicPM, funcOp.getContext());
50     if (failed(runPipeline(dynamicPM, funcOp))) return signalPassFailure();
51   }
52 };
53 
createVectorizeGmlStLoopsPass()54 std::unique_ptr<OperationPass<func::FuncOp>> createVectorizeGmlStLoopsPass() {
55   return std::make_unique<VectorizeGmlStLoopsPass>();
56 }
57 
58 }  // namespace gml_st
59 }  // namespace mlir
60