[SYCL-MLIR]: Add AlwaysInliner pass (#7481)

PR #7536 removed the MLIR inlining
pass (too greedy). This PR introduces an 'always-inline' pass in the
`cgeist` MLIR pipeline. This pass inlines only `sycl.call` call sites
with callee that has the 'alwaysinline' mlir function attribute.

The pass uses a generic SCC inliner to perform its actions and
specializes the inlining heuristic used at outline above.

Signed-off-by: Tiotto, Ettore <ettore.tiotto@intel.com>

Author: etiotto

mlir-sycl/include/mlir/Dialect/SYCL/IR/SYCLOps.td

@@ -5,21 +5,25 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+//
+// This is the operation definition file for SYCL dialect operations.
+//
+//===----------------------------------------------------------------------===//
 
 #ifndef SYCL_OPS
 #define SYCL_OPS
 
 include "mlir/IR/OpBase.td"
 include "mlir/IR/AttrTypeBase.td"
 
+include "mlir/Dialect/SYCL/IR/SYCLOpInterfaces.td"
 include "mlir/Dialect/LLVMIR/LLVMOpsInterfaces.td"
 include "mlir/Dialect/MemRef/IR/MemRefBase.td"
 include "mlir/IR/BuiltinTypeInterfaces.td"
+include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/CastInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
-include "SYCLOpInterfaces.td"
-
 ////////////////////////////////////////////////////////////////////////////////
 // TYPE DECLARATIONS
 ////////////////////////////////////////////////////////////////////////////////
@@ -77,7 +81,7 @@ class SYCL_Op<string mnemonic, list<Trait> traits = []>
 
 class SYCLMethodOpInterfaceImpl<
     string mnemonic, string type, list<string> methodNames, list<Trait> traits = []>
-        : SYCL_Op<mnemonic, !listconcat(traits, [SYCLMethodOpInterface])> {
+        : SYCL_Op<mnemonic, !listconcat(traits, [SYCLMethodOpInterface, CallOpInterface])> {
   string baseType = type;
   list<string> memberFunctionNames = methodNames;
   int arrSize = !size(memberFunctionNames);
@@ -88,10 +92,25 @@ class SYCLMethodOpInterfaceImpl<
      }};
     static ::mlir::TypeID getTypeID() { return ::mlir::sycl::}] # type # [{::getTypeID(); }
     static constexpr llvm::ArrayRef<llvm::StringLiteral> getMethodNames() { return methods; }
+
+    Operation::operand_iterator arg_operand_begin() { return (*this)->operand_begin(); }
+    Operation::operand_iterator arg_operand_end() { return (*this)->operand_end(); }
+
+    /// Return the callee of the generic SYCL call operation, this is required by
+    /// the call interface.
+    CallInterfaceCallable getCallableForCallee() {
+      return (*this)->getAttrOfType<FlatSymbolRefAttr>(getMangledFunctionNameAttrName());
+    }
+
+    /// Get the argument operands to the called function, this is required by the
+    /// call interface.
+    Operation::operand_range getArgOperands() {
+      return {arg_operand_begin(), arg_operand_end()};
+    }
   }];
 
   let extraClassDefinition = [{
-  constexpr std::array<llvm::StringLiteral, }] # arrSize # [{> $cppClass::methods;
+    constexpr std::array<llvm::StringLiteral, }] # arrSize # [{> $cppClass::methods;
   }];
 }
 
@@ -283,11 +302,10 @@ def SYCLCastOp : SYCL_Op<"cast", [DeclareOpInterfaceMethods<CastOpInterface>,
 // CALL OPERATION
 ////////////////////////////////////////////////////////////////////////////////
 
-def SYCLCallOp : SYCL_Op<"call", []> {
+def SYCLCallOp : SYCL_Op<"call", [CallOpInterface]> {
   let summary = "Generic call operation";
   let description = [{
-    This operation represent the call to any function part of the sycl's
-    namespace.
+    This operation represent a call to any function part of the sycl's namespace.
   }];
 
   let arguments = (ins
@@ -318,6 +336,23 @@ def SYCLCallOp : SYCL_Op<"call", []> {
     }]>
   ];
 
+  let extraClassDeclaration = [{
+    operand_iterator arg_operand_begin() { return operand_begin(); }
+    operand_iterator arg_operand_end() { return operand_end(); }
+
+    /// Return the callee of the generic SYCL call operation, this is required by
+    /// the call interface.
+    CallInterfaceCallable getCallableForCallee() {
+      return (*this)->getAttrOfType<FlatSymbolRefAttr>(getMangledFunctionNameAttrName());
+    }
+
+    /// Get the argument operands to the called function, this is required by the
+    /// call interface.
+    operand_range getArgOperands() {
+      return {arg_operand_begin(), arg_operand_end()};
+    }
+  }];    
+
   let assemblyFormat = [{
     `(` $Args `)` attr-dict `:` functional-type($Args, results)
   }];

mlir-sycl/include/mlir/Dialect/SYCL/Transforms/Passes.h

@@ -24,6 +24,7 @@ namespace sycl {
 // Passes
 //===----------------------------------------------------------------------===//
 
+std::unique_ptr<Pass> createAlwaysInlinePass();
 std::unique_ptr<Pass> createSYCLMethodToSYCLCallPass();
 
 //===----------------------------------------------------------------------===//

mlir-sycl/include/mlir/Dialect/SYCL/Transforms/Passes.td

@@ -11,6 +11,15 @@
 
 include "mlir/Pass/PassBase.td"
 
+def AlwaysInlinePass : Pass<"sycl-always-inline"> {
+  let summary = "Inline SYCL calls to functions marked 'alwaysinline'";
+  let description = [{
+    Replace a sycl.call operation with the body of the callee if the callee has 
+    the 'always_inline' attribute. 
+  }];  
+  let constructor = "mlir::sycl::createAlwaysInlinePass()";
+}
+
 def SYCLMethodToSYCLCall : Pass<"sycl-method-to-sycl-call", "ModuleOp"> {
   let summary = "Convert SYCLMethodOpInterface instances to SYCLCallOps";
   let description = [{

mlir-sycl/lib/Dialect/IR/SYCLOpsDialect.cpp

@@ -1,18 +1,76 @@
-// Copyright (C) Codeplay Software Limited
-
-//===--- SYCLOpsDialect.cpp -----------------------------------------------===//
+//===--- SYCLOpsDialect.cpp - SYCL Dialect registration in MLIR -----------===//
 //
 // MLIR-SYCL is under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+//
+// This file implements the dialect for the SYCL IR.
+//
+//===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SYCL/IR/SYCLOpsDialect.h"
-
 #include "mlir/Dialect/SYCL/IR/SYCLOps.h"
 #include "mlir/Dialect/SYCL/IR/SYCLOpsAlias.h"
 #include "mlir/Dialect/SYCL/IR/SYCLOpsTypes.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "mlir/Transforms/InliningUtils.h"
+
+//===----------------------------------------------------------------------===//
+// SYCL Dialect Interfaces
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// This class defines the interface for inlining SYCL operations.
+class SYCLInlinerInterface : public mlir::DialectInlinerInterface {
+public:
+  using DialectInlinerInterface::DialectInlinerInterface;
+
+  //===--------------------------------------------------------------------===//
+  // Analysis Hooks
+  //===--------------------------------------------------------------------===//
+
+  /// This hook checks whether is legal to inline the \p Callable operation and
+  /// replace the \p Call operation with it. For the SYCL dialect we want to
+  /// allow inlining only SYCLCallOp operations.
+  bool isLegalToInline(mlir::Operation *Call, mlir::Operation *Callable,
+                       bool WouldBeCloned) const final {
+    return mlir::isa<mlir::sycl::SYCLCallOp>(Call);
+  }
+
+  /// This hook checks whether is legal to inline the \p Op operation into the
+  /// \p Dest region. All operations in the SYCL dialect are legal to inline.
+  bool isLegalToInline(mlir::Operation *Op, mlir::Region *Dest,
+                       bool WouldBeCloned,
+                       mlir::BlockAndValueMapping &ValueMapping) const final {
+    return true;
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Transformation Hooks
+  //===--------------------------------------------------------------------===//
+
+  /// Attempts to materialize a conversion for a type mismatch between a call
+  /// from the SYCL dialect, and a callable region. This method should generate
+  /// an operation that takes \p Input as the only operand, and produces a
+  /// single result of \p ResultType. If a conversion cannot be generated,
+  /// nullptr should be returned.
+  mlir::Operation *
+  materializeCallConversion(mlir::OpBuilder &Builder, mlir::Value Input,
+                            mlir::Type ResultType,
+                            mlir::Location ConversionLoc) const final {
+    return Builder.create<mlir::sycl::SYCLCastOp>(ConversionLoc, ResultType,
+                                                  Input);
+  }
+};
+
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// SYCL Dialect
+//===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
@@ -37,6 +95,7 @@ void mlir::sycl::SYCLDialect::initialize() {
       >();
 
   mlir::Dialect::addInterfaces<SYCLOpAsmInterface>();
+  mlir::Dialect::addInterfaces<SYCLInlinerInterface>();
 }
 
 llvm::Optional<llvm::StringRef>

mlir-sycl/lib/Transforms/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_mlir_library(MLIRSYCLTransforms
+  Inliner.cpp
   SYCLMethodToSYCLCall.cpp
 
   ADDITIONAL_HEADER_DIRS
@@ -9,6 +10,9 @@ add_mlir_library(MLIRSYCLTransforms
 
   LINK_LIBS PUBLIC
   MLIRPass
+  MLIRTransforms  
   MLIRTransformUtils
   MLIRSYCLDialect
+  MLIRFuncDialect
+  MLIRGPUOps
   )