diff --git a/mlir/include/mlir/Dialect/CMakeLists.txt b/mlir/include/mlir/Dialect/CMakeLists.txt --- a/mlir/include/mlir/Dialect/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/CMakeLists.txt @@ -14,6 +14,7 @@ add_subdirectory(Linalg) add_subdirectory(LLVMIR) add_subdirectory(MemRef) +add_subdirectory(MLProgram) add_subdirectory(OpenACC) add_subdirectory(OpenMP) add_subdirectory(PDL) diff --git a/mlir/include/mlir/Dialect/MLProgram/CMakeLists.txt b/mlir/include/mlir/Dialect/MLProgram/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/MLProgram/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/mlir/include/mlir/Dialect/MLProgram/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/MLProgram/IR/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/MLProgram/IR/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS MLProgramOps.td) +add_mlir_dialect(MLProgramOps ml_program) +add_mlir_doc(MLProgramOps MLProgramOps Dialects/ -gen-dialect-doc) diff --git a/mlir/include/mlir/Dialect/MLProgram/IR/MLProgram.h b/mlir/include/mlir/Dialect/MLProgram/IR/MLProgram.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/MLProgram/IR/MLProgram.h @@ -0,0 +1,30 @@ +//===- MLProgram.h - MLProgram dialect ----------------------------*- C++-*-==// +// +// Part of the LLVM Project, 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 +// +//===----------------------------------------------------------------------===// +#ifndef MLIR_DIALECT_MLPROGRAM_IR_MLPROGRAM_H_ +#define MLIR_DIALECT_MLPROGRAM_IR_MLPROGRAM_H_ + +#include "mlir/IR/Dialect.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/OpImplementation.h" +#include "mlir/IR/RegionKindInterface.h" +#include "mlir/IR/SymbolTable.h" + +//===----------------------------------------------------------------------===// +// MLProgramDialect +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/MLProgram/IR/MLProgramOpsDialect.h.inc" + +//===----------------------------------------------------------------------===// +// MLProgram Dialect Operations +//===----------------------------------------------------------------------===// + +#define GET_OP_CLASSES +#include "mlir/Dialect/MLProgram/IR/MLProgramOps.h.inc" + +#endif // MLIR_DIALECT_MLPROGRAM_IR_MLPROGRAM_H_ diff --git a/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramBase.td b/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramBase.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramBase.td @@ -0,0 +1,29 @@ +//===- MLProgramBase.td - Base defs for ml_program dialect --*- tablegen -*-==// +// +// Part of the LLVM Project, 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 +// +//===----------------------------------------------------------------------===// + +#ifndef MLPROGRAM_BASE +#define MLPROGRAM_BASE + +include "mlir/IR/OpBase.td" + +def MLProgram_Dialect : Dialect { + let name = "ml_program"; + let cppNamespace = "::mlir::ml_program"; + let description = [{ + The MLProgram dialect contains structural operations and types for + defining a compiled Machine-Learning program, as created from common + ML frameworks, such as TensorFlow, PyTorch, JAX, etc. It does not itself + define computation ops common to such frameworks but establishes a common + programming model for establishing modules, functions, globals and + memory model components appropriate for such an abstract level of detail. + }]; + + let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed; +} + +#endif // MLPROGRAM_BASE diff --git a/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td b/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td @@ -0,0 +1,54 @@ +//===- MLProgramOps.td - Structural ML Program Ops ---------*- tablegen -*-===// +// +// Part of the LLVM Project, 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 +// +//===----------------------------------------------------------------------===// + +#ifndef MLPROGRAM_OPS +#define MLPROGRAM_OPS + +include "mlir/Dialect/MLProgram/IR/MLProgramBase.td" +include "mlir/IR/RegionKindInterface.td" +include "mlir/IR/SymbolInterfaces.td" + +class MLProgram_Op traits = []> : + Op; + +//===----------------------------------------------------------------------===// +// ModuleOp +//===----------------------------------------------------------------------===// + +def ModuleOp : MLProgram_Op<"module", [ + IsolatedFromAbove, NoRegionArguments, SymbolTable, Symbol + ] # GraphRegionNoTerminator.traits> { + let summary = "A self contained ML program module"; + let description = [{ + A `module` represents a container for top-level ML program ops. Modules + have an optional symbol name which can be used to refer to them in + operations and differentiate them in collections of modules. + }]; + + let arguments = (ins OptionalAttr:$sym_name, + OptionalAttr:$sym_visibility); + let regions = (region SizedRegion<1>:$body); + + let assemblyFormat = "($sym_name^)? attr-dict-with-keyword $body"; + let builders = [OpBuilder<(ins CArg<"Optional", "{}">:$name)>]; + + // We need to ensure the block inside the region is properly terminated; + // the auto-generated builders do not guarantee that. + let skipDefaultBuilders = 1; + + let extraClassDeclaration = [{ + //===------------------------------------------------------------------===// + // SymbolOpInterface Methods + //===------------------------------------------------------------------===// + + /// A ModuleOp may optionally define a symbol. + bool isOptionalSymbol() { return true; } + }]; +} + +#endif // MLPROGRAM_OPS diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h --- a/mlir/include/mlir/InitAllDialects.h +++ b/mlir/include/mlir/InitAllDialects.h @@ -30,6 +30,7 @@ #include "mlir/Dialect/LLVMIR/NVVMDialect.h" #include "mlir/Dialect/LLVMIR/ROCDLDialect.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MLProgram/IR/MLProgram.h" #include "mlir/Dialect/Math/IR/Math.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/OpenACC/OpenACC.h" @@ -71,6 +72,7 @@ linalg::LinalgDialect, math::MathDialect, memref::MemRefDialect, + ml_program::MLProgramDialect, scf::SCFDialect, omp::OpenMPDialect, pdl::PDLDialect, diff --git a/mlir/lib/Dialect/CMakeLists.txt b/mlir/lib/Dialect/CMakeLists.txt --- a/mlir/lib/Dialect/CMakeLists.txt +++ b/mlir/lib/Dialect/CMakeLists.txt @@ -14,6 +14,7 @@ add_subdirectory(LLVMIR) add_subdirectory(Math) add_subdirectory(MemRef) +add_subdirectory(MLProgram) add_subdirectory(OpenACC) add_subdirectory(OpenMP) add_subdirectory(PDL) diff --git a/mlir/lib/Dialect/MLProgram/CMakeLists.txt b/mlir/lib/Dialect/MLProgram/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/MLProgram/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt b/mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt @@ -0,0 +1,15 @@ +add_mlir_dialect_library(MLIRMLProgram + MLProgramOps.cpp + MLProgramDialect.cpp + + ADDITIONAL_HEADER_DIRS + ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MLProgram + + DEPENDS + MLIRMLProgramOpsIncGen + + LINK_LIBS PUBLIC + MLIRDialect + MLIRInferTypeOpInterface + MLIRIR + ) diff --git a/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp b/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp @@ -0,0 +1,21 @@ +//===- MLProgramDialect.cpp - MLIR MLProgram dialect implementation -------===// +// +// Part of the LLVM Project, 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/MLProgram/IR/MLProgram.h" + +using namespace mlir; +using namespace mlir::ml_program; + +#include "mlir/Dialect/MLProgram/IR/MLProgramOpsDialect.cpp.inc" + +void ml_program::MLProgramDialect::initialize() { + addOperations< +#define GET_OP_LIST +#include "mlir/Dialect/MLProgram/IR/MLProgramOps.cpp.inc" + >(); +} diff --git a/mlir/lib/Dialect/MLProgram/IR/MLProgramOps.cpp b/mlir/lib/Dialect/MLProgram/IR/MLProgramOps.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/MLProgram/IR/MLProgramOps.cpp @@ -0,0 +1,33 @@ +//===- MLProgramOps.cpp - MLIR MLProgram dialect ops implementation -------===// +// +// Part of the LLVM Project, 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/MLProgram/IR/MLProgram.h" +#include "mlir/IR/Builders.h" + +using namespace mlir; +using namespace mlir::ml_program; + +//===----------------------------------------------------------------------===// +// TableGen'd op method definitions +//===----------------------------------------------------------------------===// + +#define GET_OP_CLASSES +#include "mlir/Dialect/MLProgram/IR/MLProgramOps.cpp.inc" + +//===----------------------------------------------------------------------===// +// ModuleOp +//===----------------------------------------------------------------------===// + +void ModuleOp::build(OpBuilder &builder, OperationState &state, + Optional name) { + state.addRegion()->emplaceBlock(); + if (name) { + state.attributes.push_back(builder.getNamedAttr( + mlir::SymbolTable::getSymbolAttrName(), builder.getStringAttr(*name))); + } +} diff --git a/mlir/test/Dialect/MLProgram/ops.mlir b/mlir/test/Dialect/MLProgram/ops.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/MLProgram/ops.mlir @@ -0,0 +1,10 @@ +// RUN: mlir-opt %s | mlir-opt | FileCheck %s +// RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s + +// CHECK-LABEL: ml_program.module { +ml_program.module { +} + +// CHECK-LABEL: ml_program.module @foobar { +ml_program.module @foobar { +} diff --git a/mlir/test/mlir-opt/commandline.mlir b/mlir/test/mlir-opt/commandline.mlir --- a/mlir/test/mlir-opt/commandline.mlir +++ b/mlir/test/mlir-opt/commandline.mlir @@ -18,6 +18,7 @@ // CHECK-NEXT: llvm // CHECK-NEXT: math // CHECK-NEXT: memref +// CHECK-NEXT: ml_program // CHECK-NEXT: nvvm // CHECK-NEXT: omp // CHECK-NEXT: pdl