Make ConvertToSPIRV passes OperationPass

[jackalcooper]

Summary

Make passes converting ops from other dialects to spirv OperationPass, so that downstream compiler could put them in a proper nested pass manager to lower device code only.

[antiagainst]

Awesome, thanks!

mlir/include/mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h

16

This is not for SPIR-V. Could you split it out as a separate patch? Thanks!

[jackalcooper]

mlir/include/mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h

16

This include pulls in the declaration for OperationPass<> used below.

[jackalcooper]

mlir/include/mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h

16

Oh, I can use fw declaration.

[antiagainst]

mlir/include/mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h

16

Sorry, the comment is misplaced. This is actually for L32, which is createConvertControlFlowToLLVMPass().

mlir/include/mlir/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.h

 #ifndef MLIR_CONVERSION_ARITHMETICTOSPIRV_ARITHMETICTOSPIRV_H
 #define MLIR_CONVERSION_ARITHMETICTOSPIRV_ARITHMETICTOSPIRV_H
 
+#include "mlir/Pass/Pass.h"
 #include <memory>
 
 namespace mlir {

 void populateArithmeticToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                        RewritePatternSet &patterns);
 
-std::unique_ptr<Pass> createConvertArithmeticToSPIRVPass();
+std::unique_ptr<OperationPass<>> createConvertArithmeticToSPIRVPass();
 } // namespace arith
 } // namespace mlir
 

mlir/include/mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h

 #ifndef MLIR_CONVERSION_CONTROLFLOWTOLLVM_CONTROLFLOWTOLLVM_H
 #define MLIR_CONVERSION_CONTROLFLOWTOLLVM_CONTROLFLOWTOLLVM_H
 
+#include "mlir/Pass/Pass.h"
 #include <memory>
 
 namespace mlir {

                                                  RewritePatternSet &patterns);
 
 /// Creates a pass to convert the ControlFlow dialect into the LLVMIR dialect.
-std::unique_ptr<Pass> createConvertControlFlowToLLVMPass();
+std::unique_ptr<OperationPass<>> createConvertControlFlowToLLVMPass();
 } // namespace cf
 } // namespace mlir
 

mlir/include/mlir/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRVPass.h

 class ModuleOp;
 
 /// Creates a pass to convert ControlFlow ops to SPIR-V ops.
-std::unique_ptr<OperationPass<ModuleOp>> createConvertControlFlowToSPIRVPass();
+std::unique_ptr<OperationPass<>> createConvertControlFlowToSPIRVPass();
 
 } // namespace mlir
 

mlir/include/mlir/Conversion/FuncToSPIRV/FuncToSPIRVPass.h

 class ModuleOp;
 
 /// Creates a pass to convert Func ops to SPIR-V ops.
-std::unique_ptr<OperationPass<ModuleOp>> createConvertFuncToSPIRVPass();
+std::unique_ptr<OperationPass<>> createConvertFuncToSPIRVPass();
 
 } // namespace mlir
 

mlir/include/mlir/Conversion/MathToSPIRV/MathToSPIRVPass.h

 class ModuleOp;
 
 /// Creates a pass to convert Math ops to SPIR-V ops.
-std::unique_ptr<OperationPass<ModuleOp>> createConvertMathToSPIRVPass();
+std::unique_ptr<OperationPass<>> createConvertMathToSPIRVPass();
 
 } // namespace mlir
 

mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h

 std::unique_ptr<OperationPass<>> createMapMemRefStorageClassPass();
 
 /// Creates a pass to convert MemRef ops to SPIR-V ops.
-std::unique_ptr<OperationPass<ModuleOp>> createConvertMemRefToSPIRVPass();
+std::unique_ptr<OperationPass<>> createConvertMemRefToSPIRVPass();
 
 } // namespace mlir
 

mlir/include/mlir/Conversion/Passes.td

 // ArithmeticToSPIRV
 //===----------------------------------------------------------------------===//
 
-def ConvertArithmeticToSPIRV : Pass<"convert-arith-to-spirv", "ModuleOp"> {
+def ConvertArithmeticToSPIRV : Pass<"convert-arith-to-spirv"> {
   let summary = "Convert Arithmetic dialect to SPIR-V dialect";
   let constructor = "mlir::arith::createConvertArithmeticToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];

 // ControlFlowToLLVM
 //===----------------------------------------------------------------------===//
 
-def ConvertControlFlowToLLVM : Pass<"convert-cf-to-llvm", "ModuleOp"> {
+def ConvertControlFlowToLLVM : Pass<"convert-cf-to-llvm"> {
   let summary = "Convert ControlFlow operations to the LLVM dialect";
   let description = [{
     Convert ControlFlow operations into LLVM IR dialect operations.

 // ControlFlowToSPIRV
 //===----------------------------------------------------------------------===//
 
-def ConvertControlFlowToSPIRV : Pass<"convert-cf-to-spirv", "ModuleOp"> {
+def ConvertControlFlowToSPIRV : Pass<"convert-cf-to-spirv"> {
   let summary = "Convert ControlFlow dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertControlFlowToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];

 // FuncToSPIRV
 //===----------------------------------------------------------------------===//
 
-def ConvertFuncToSPIRV : Pass<"convert-func-to-spirv", "ModuleOp"> {
+def ConvertFuncToSPIRV : Pass<"convert-func-to-spirv"> {
   let summary = "Convert Func dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertFuncToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];

 // MathToSPIRV
 //===----------------------------------------------------------------------===//
 
-def ConvertMathToSPIRV : Pass<"convert-math-to-spirv", "ModuleOp"> {
+def ConvertMathToSPIRV : Pass<"convert-math-to-spirv"> {
   let summary = "Convert Math dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertMathToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];

   ];
 }
 
-def ConvertMemRefToSPIRV : Pass<"convert-memref-to-spirv", "ModuleOp"> {
+def ConvertMemRefToSPIRV : Pass<"convert-memref-to-spirv"> {
   let summary = "Convert MemRef dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertMemRefToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];

 // SCFToSPIRV
 //===----------------------------------------------------------------------===//
 
-def SCFToSPIRV : Pass<"convert-scf-to-spirv", "ModuleOp"> {
+def SCFToSPIRV : Pass<"convert-scf-to-spirv"> {
   let summary = "Convert SCF dialect to SPIR-V dialect.";
   let description = [{
     This pass converts SCF ops into SPIR-V structured control flow ops.

 // TensorToSPIRV
 //===----------------------------------------------------------------------===//
 
-def ConvertTensorToSPIRV : Pass<"convert-tensor-to-spirv", "ModuleOp"> {
+def ConvertTensorToSPIRV : Pass<"convert-tensor-to-spirv"> {
   let summary = "Convert Tensor dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertTensorToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];

 // VectorToSPIRV
 //===----------------------------------------------------------------------===//
 
-def ConvertVectorToSPIRV : Pass<"convert-vector-to-spirv", "ModuleOp"> {
+def ConvertVectorToSPIRV : Pass<"convert-vector-to-spirv"> {
   let summary = "Convert Vector dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertVectorToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];

mlir/include/mlir/Conversion/SCFToSPIRV/SCFToSPIRVPass.h

 class ModuleOp;
 
 /// Creates a pass to convert SCF ops into SPIR-V ops.
-std::unique_ptr<OperationPass<ModuleOp>> createConvertSCFToSPIRVPass();
+std::unique_ptr<OperationPass<>> createConvertSCFToSPIRVPass();
 
 } // namespace mlir
 

mlir/include/mlir/Conversion/TensorToSPIRV/TensorToSPIRVPass.h

 class ModuleOp;
 
 /// Creates a pass to convert Tensor ops to SPIR-V ops.
-std::unique_ptr<OperationPass<ModuleOp>> createConvertTensorToSPIRVPass();
+std::unique_ptr<OperationPass<>> createConvertTensorToSPIRVPass();
 
 } // namespace mlir
 

mlir/include/mlir/Conversion/VectorToSPIRV/VectorToSPIRVPass.h

 class ModuleOp;
 
 /// Creates a pass to convert Vector Ops to SPIR-V ops.
-std::unique_ptr<OperationPass<ModuleOp>> createConvertVectorToSPIRVPass();
+std::unique_ptr<OperationPass<>> createConvertVectorToSPIRVPass();
 
 } // namespace mlir
 

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

 struct ConvertArithmeticToSPIRVPass
     : public ConvertArithmeticToSPIRVBase<ConvertArithmeticToSPIRVPass> {
   void runOnOperation() override {
-    auto module = getOperation();
-    auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+    Operation *op = getOperation();
+    auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
     auto target = SPIRVConversionTarget::get(targetAttr);
 
     SPIRVTypeConverter::Options options;

     RewritePatternSet patterns(&getContext());
     arith::populateArithmeticToSPIRVPatterns(typeConverter, patterns);
 
-    if (failed(applyPartialConversion(module, *target, std::move(patterns))))
+    if (failed(applyPartialConversion(op, *target, std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
-std::unique_ptr<Pass> mlir::arith::createConvertArithmeticToSPIRVPass() {
+std::unique_ptr<OperationPass<>>
+mlir::arith::createConvertArithmeticToSPIRVPass() {
   return std::make_unique<ConvertArithmeticToSPIRVPass>();
 }

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

 };
 } // namespace
 
-std::unique_ptr<Pass> mlir::cf::createConvertControlFlowToLLVMPass() {
+std::unique_ptr<OperationPass<>>
+mlir::cf::createConvertControlFlowToLLVMPass() {
   return std::make_unique<ConvertControlFlowToLLVM>();
 }

mlir/lib/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRVPass.cpp

 
 void ConvertControlFlowToSPIRVPass::runOnOperation() {
   MLIRContext *context = &getContext();
-  ModuleOp module = getOperation();
+  Operation *op = getOperation();
 
-  auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+  auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
   std::unique_ptr<ConversionTarget> target =
       SPIRVConversionTarget::get(targetAttr);
 

   RewritePatternSet patterns(context);
   cf::populateControlFlowToSPIRVPatterns(typeConverter, patterns);
 
-  if (failed(applyPartialConversion(module, *target, std::move(patterns))))
+  if (failed(applyPartialConversion(op, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
-std::unique_ptr<OperationPass<ModuleOp>>
-mlir::createConvertControlFlowToSPIRVPass() {
+std::unique_ptr<OperationPass<>> mlir::createConvertControlFlowToSPIRVPass() {
   return std::make_unique<ConvertControlFlowToSPIRVPass>();
 }

mlir/lib/Conversion/FuncToSPIRV/FuncToSPIRVPass.cpp

 
 void ConvertFuncToSPIRVPass::runOnOperation() {
   MLIRContext *context = &getContext();
-  ModuleOp module = getOperation();
+  Operation *op = getOperation();
 
-  auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+  auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
   std::unique_ptr<ConversionTarget> target =
       SPIRVConversionTarget::get(targetAttr);
 

   populateFuncToSPIRVPatterns(typeConverter, patterns);
   populateBuiltinFuncToSPIRVPatterns(typeConverter, patterns);
 
-  if (failed(applyPartialConversion(module, *target, std::move(patterns))))
+  if (failed(applyPartialConversion(op, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
-std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertFuncToSPIRVPass() {
+std::unique_ptr<OperationPass<>> mlir::createConvertFuncToSPIRVPass() {
   return std::make_unique<ConvertFuncToSPIRVPass>();
 }

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp

     gpuModules.push_back(builder.clone(*moduleOp.getOperation()));
   });
 
-  // Map MemRef memory space to SPIR-V sotrage class first if requested.
+  // Map MemRef memory space to SPIR-V storage class first if requested.
   if (mapMemorySpace) {
     std::unique_ptr<ConversionTarget> target =
         spirv::getMemorySpaceToStorageClassTarget(*context);

mlir/lib/Conversion/MathToSPIRV/MathToSPIRVPass.cpp

 
 void ConvertMathToSPIRVPass::runOnOperation() {
   MLIRContext *context = &getContext();
-  ModuleOp module = getOperation();
+  Operation *op = getOperation();
 
-  auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+  auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
   std::unique_ptr<ConversionTarget> target =
       SPIRVConversionTarget::get(targetAttr);
 

   RewritePatternSet patterns(context);
   populateMathToSPIRVPatterns(typeConverter, patterns);
 
-  if (failed(applyPartialConversion(module, *target, std::move(patterns))))
+  if (failed(applyPartialConversion(op, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
-std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertMathToSPIRVPass() {
+std::unique_ptr<OperationPass<>> mlir::createConvertMathToSPIRVPass() {
   return std::make_unique<ConvertMathToSPIRVPass>();
 }

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.cpp

 
 void ConvertMemRefToSPIRVPass::runOnOperation() {
   MLIRContext *context = &getContext();
-  ModuleOp module = getOperation();
+  Operation *op = getOperation();
 
-  auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+  auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
   std::unique_ptr<ConversionTarget> target =
       SPIRVConversionTarget::get(targetAttr);
 

   RewritePatternSet patterns(context);
   populateMemRefToSPIRVPatterns(typeConverter, patterns);
 
-  if (failed(applyPartialConversion(module, *target, std::move(patterns))))
+  if (failed(applyPartialConversion(op, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
-std::unique_ptr<OperationPass<ModuleOp>>
-mlir::createConvertMemRefToSPIRVPass() {
+std::unique_ptr<OperationPass<>> mlir::createConvertMemRefToSPIRVPass() {
   return std::make_unique<ConvertMemRefToSPIRVPass>();
 }

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRVPass.cpp

 
 void SCFToSPIRVPass::runOnOperation() {
   MLIRContext *context = &getContext();
-  ModuleOp module = getOperation();
+  Operation *op = getOperation();
 
-  auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+  auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
   std::unique_ptr<ConversionTarget> target =
       SPIRVConversionTarget::get(targetAttr);
 

   populateMemRefToSPIRVPatterns(typeConverter, patterns);
   populateBuiltinFuncToSPIRVPatterns(typeConverter, patterns);
 
-  if (failed(applyPartialConversion(module, *target, std::move(patterns))))
+  if (failed(applyPartialConversion(op, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
-std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertSCFToSPIRVPass() {
+std::unique_ptr<OperationPass<>> mlir::createConvertSCFToSPIRVPass() {
   return std::make_unique<SCFToSPIRVPass>();
 }

mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRVPass.cpp

     : public ConvertTensorToSPIRVBase<ConvertTensorToSPIRVPass> {
   void runOnOperation() override {
     MLIRContext *context = &getContext();
-    ModuleOp module = getOperation();
+    Operation *op = getOperation();
 
-    auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+    auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
     std::unique_ptr<ConversionTarget> target =
         SPIRVConversionTarget::get(targetAttr);
 

                                   patterns);
     populateBuiltinFuncToSPIRVPatterns(typeConverter, patterns);
 
-    if (failed(applyPartialConversion(module, *target, std::move(patterns))))
+    if (failed(applyPartialConversion(op, *target, std::move(patterns))))
       return signalPassFailure();
   }
 };
 } // namespace
 
-std::unique_ptr<OperationPass<ModuleOp>>
-mlir::createConvertTensorToSPIRVPass() {
+std::unique_ptr<OperationPass<>> mlir::createConvertTensorToSPIRVPass() {
   return std::make_unique<ConvertTensorToSPIRVPass>();
 }

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRVPass.cpp

 
 void ConvertVectorToSPIRVPass::runOnOperation() {
   MLIRContext *context = &getContext();
-  ModuleOp module = getOperation();
+  Operation *op = getOperation();
 
-  auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+  auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
   std::unique_ptr<ConversionTarget> target =
       SPIRVConversionTarget::get(targetAttr);
 

   RewritePatternSet patterns(context);
   populateVectorToSPIRVPatterns(typeConverter, patterns);
 
-  if (failed(applyPartialConversion(module, *target, std::move(patterns))))
+  if (failed(applyPartialConversion(op, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
-std::unique_ptr<OperationPass<ModuleOp>>
-mlir::createConvertVectorToSPIRVPass() {
+std::unique_ptr<OperationPass<>> mlir::createConvertVectorToSPIRVPass() {
   return std::make_unique<ConvertVectorToSPIRVPass>();
 }