From 0efd79e0d28e54a297686317d36dbc371575a22f Mon Sep 17 00:00:00 2001 From: Stella Laurenzo Date: Fri, 15 Mar 2024 15:30:25 -0700 Subject: [PATCH 1/2] [mlir] Make the ml_program dialect allow all of its operations to be inlined. --- mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt | 1 + mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp | 14 +++++++++++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt b/mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt index 90c9c9aa78750..725bb5fd9da9e 100644 --- a/mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt @@ -15,5 +15,6 @@ add_mlir_dialect_library(MLIRMLProgramDialect MLIRControlFlowInterfaces MLIRFunctionInterfaces MLIRInferTypeOpInterface + MLIRTransforms MLIRIR ) diff --git a/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp b/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp index 1a8fe208d4099..0b186a0d072aa 100644 --- a/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp +++ b/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp @@ -8,6 +8,7 @@ #include "mlir/Dialect/MLProgram/IR/MLProgram.h" #include "mlir/IR/DialectImplementation.h" +#include "mlir/Transforms/InliningUtils.h" #include "llvm/ADT/TypeSwitch.h" using namespace mlir; @@ -24,6 +25,17 @@ using namespace mlir::ml_program; #include "mlir/Dialect/MLProgram/IR/MLProgramTypes.cpp.inc" namespace { + +struct MLProgramInlinerInterface : public DialectInlinerInterface { + using DialectInlinerInterface::DialectInlinerInterface; + + bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const { + // We have no specific opinion on whether ops defined in this dialect should + // be inlined. + return true; + } +}; + struct MLProgramOpAsmDialectInterface : public OpAsmDialectInterface { using OpAsmDialectInterface::OpAsmDialectInterface; @@ -53,5 +65,5 @@ void ml_program::MLProgramDialect::initialize() { #include "mlir/Dialect/MLProgram/IR/MLProgramOps.cpp.inc" >(); - addInterfaces(); + addInterfaces(); } From 96ba583ce9048769bf255b1cb6af30f0c21d0a9d Mon Sep 17 00:00:00 2001 From: Stella Laurenzo Date: Fri, 15 Mar 2024 16:08:11 -0700 Subject: [PATCH 2/2] Add test --- .../Dialect/MLProgram/IR/MLProgramDialect.cpp | 3 ++- mlir/test/Dialect/MLProgram/inlining.mlir | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) create mode 100644 mlir/test/Dialect/MLProgram/inlining.mlir diff --git a/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp b/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp index 0b186a0d072aa..bda1032ed9884 100644 --- a/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp +++ b/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp @@ -29,7 +29,8 @@ namespace { struct MLProgramInlinerInterface : public DialectInlinerInterface { using DialectInlinerInterface::DialectInlinerInterface; - bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const { + bool isLegalToInline(Operation *, Region *, bool, + IRMapping &) const override { // We have no specific opinion on whether ops defined in this dialect should // be inlined. return true; diff --git a/mlir/test/Dialect/MLProgram/inlining.mlir b/mlir/test/Dialect/MLProgram/inlining.mlir new file mode 100644 index 0000000000000..ac3677e1b9288 --- /dev/null +++ b/mlir/test/Dialect/MLProgram/inlining.mlir @@ -0,0 +1,19 @@ +// RUN: mlir-opt %s -inline | FileCheck %s + +// Verifies that regions with operations from the ml_program dialect can +// be inlined. + +ml_program.global private @global(dense<4> : tensor<4xi32>) : tensor<4xi32> + +// CHECK: @inline_into +func.func @inline_into() -> tensor<4xi32> { + // CHECK-NOT: @inline_from + // CHECK: ml_program.global_load_const + %0 = call @inline_from() : () -> tensor<4xi32> + return %0 : tensor<4xi32> +} + +func.func @inline_from() -> tensor<4xi32> { + %0 = ml_program.global_load_const @global : tensor<4xi32> + return %0 : tensor<4xi32> +}