[mlir] Introduce ml_program dialect.

[stellaraccident]

Add func/subgraph ops.

[Mogball]

mlir/include/mlir/IR/FunctionInterfaces.td

89 ↗

(On Diff #421396)

Could the changes in this file be split to a separate patch? (Or just pushed as an NFC...)

[stellaraccident]

mlir/include/mlir/IR/FunctionInterfaces.td

89 ↗

(On Diff #421396)

Thanks - meant to do that but forgot. Will split before landing.

[silvas]

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td

170

what does "terminator" mean for a graph region? Can you just remove it (it also says "terminates a subgraph" below too.

174

nit: remove "function" from this, to avoid confusion. or is the idea that subgraph and func are both "functions"?

mlir/lib/Dialect/MLProgram/IR/MLProgramOps.cpp

72

nit: not sure if "returns" is the most consistent nomenclature here. Maybe "has N results"?

76

same comment about "return"

[mehdi_amini]

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td

170

In tensorflow, TFG Functions are graph region but the FuncOp verifier enforces that there is a terminator still (the FuncOp return)

[lattner]

Nice!

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td

40

Grammar nit: it cannot represented nested symbols.

65

plz fit to 80 cols here and below

79

Are these provided by FunctionOpInterface now?

170

I'm not sure what Sean means here: a terminator is the same in a graph region as a CFG region: it is the last operation in the block. This is important even though "order of execution" isn't generally a thing in graph regions. You still need a canonical place to put it, just like how LLVM puts PHIs at the top of a block.

[jpienaar]

Nice :-)

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td

15

If you added the infer return type interface this should result in some simpler builders be generated too I believe (just looking at signatures)

32

s/An f/F/ ?

174

OOC why `` around subgraph (here and above)? It doesn't variable or accessor and subgraph is well established term in graph theory (I find it common enough to be akin to function in needing special markers)

203

(for later discussion) I will mention that we found exact matching quite cumbersome during progressive lowering and optimizations. For one it doesnt match the ML frameworks we were working with and so we had casts that weren't needed by the system (and in this case it means one would probably only partially use ML program dialect and only post rounds of high level optimizations rather than introduce it early). One ends up inserting numerous casts to appease the type system, these may then obscure other pattern matches (while the ops interacting are perfectly happy with compatible types) and one ends up invoking shape inference more times than needed to identify and nuke unneeded casts (while some casts may have performance impact due to device assignment in distributed workload, so it isn't always simple removal).

mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp

2

Nit: MLIR feels redundant here.

mlir/lib/Dialect/MLProgram/IR/MLProgramOps.cpp

99

(I'd almost be tempted to do llvm::enumerate combined with zip here, but would need to try to see how readable)

[silvas]

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td

170

Oh, I think I must have been confusing graph regions with the "region doesn't need to have a terminator" support that Mehdi added for ModuleOp. Sorry for the noise.

[silvas]

Can you add a pass / reusable pattern that converts ml_program.func to func.func and the reverse? Otherwise I feel like we are going to roll that many times downstream just for our initial push in IREE (can be a separate patch of course).

[stellaraccident]

Hi folks - thanks for the reviews and sorry to leave this idling. I came down with covid last week and am just getting back online. Will address comments and move this forward soon.

[stellaraccident]

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td

15

Really? Even though none of these ops have return types?

65

Oops. Thanks. Used to formatters/linters :/

79

No, they are just declarations on that interface.

203

I thought about this and opted to do this for now and discuss further. Completely open to broadening in a followup.

mlir/include/mlir/IR/FunctionInterfaces.td

89 ↗

(On Diff #421396)

Split into https://reviews.llvm.org/D123757

mlir/lib/Dialect/MLProgram/IR/MLProgramOps.cpp

72

Thanks - copypasta. Updated here and below.

99

I got this from func.return so we should probably update there if we come up with something more readable. I don't have an opinion. Per your other comment, we're likely going to re-evaluate the verification behavior anyway, which will mean rewriting this.

[stellaraccident]

Rebase and comments.

[stellaraccident]

Fix test.

[stellaraccident]

Rebase

[stellaraccident]

Rebase.

mlir/include/mlir/Dialect/CMakeLists.txt

 add_subdirectory(DLTI)
 add_subdirectory(EmitC)
 add_subdirectory(Func)
 add_subdirectory(GPU)
 add_subdirectory(Math)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(MemRef)
+add_subdirectory(MLProgram)
 add_subdirectory(OpenACC)
 add_subdirectory(OpenMP)
 add_subdirectory(PDL)
 add_subdirectory(PDLInterp)
 add_subdirectory(Quant)
 add_subdirectory(SCF)
 add_subdirectory(Shape)
 add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)
 add_subdirectory(Tensor)
 add_subdirectory(Tosa)
 add_subdirectory(Vector)
 add_subdirectory(X86Vector)

mlir/include/mlir/Dialect/MLProgram/CMakeLists.txt

+add_subdirectory(IR)

mlir/include/mlir/Dialect/MLProgram/IR/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS MLProgramOps.td)
+add_mlir_dialect(MLProgramOps ml_program)
+add_mlir_doc(MLProgramOps MLProgramOps Dialects/ -gen-dialect-doc)

mlir/include/mlir/Dialect/MLProgram/IR/MLProgram.h

+//===- 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/FunctionInterfaces.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/RegionKindInterface.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Interfaces/CallInterfaces.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/SideEffectInterfaces.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_

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramBase.td

+//===- 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.
+
+    This dialect is under active development, and while stability is an
+    eventual goal, it is not guaranteed at this juncture. Given the early state,
+    it is recommended to inquire further prior to using this dialect.
+  }];
+
+  let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed;
+}
+
+#endif // MLPROGRAM_BASE

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td

+//===- 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/Interfaces/CallInterfaces.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/FunctionInterfaces.td"
+include "mlir/IR/RegionKindInterface.td"
+include "mlir/IR/SymbolInterfaces.td"
+
+class MLProgram_Op<string mnemonic, list<Trait> traits = []> :
+    Op<MLProgram_Dialect, mnemonic, traits>;
+
+//===----------------------------------------------------------------------===//
+// FuncOp
+//===----------------------------------------------------------------------===//
+
+def MLProgram_FuncOp : MLProgram_Op<"func", [
+    CallableOpInterface, FunctionOpInterface, IsolatedFromAbove,
+    RegionKindInterface, Symbol
+  ]> {
+  let summary = "Function containing a single `SSACFG` region";
+  let description = [{
+    This simple function container represents callables in an ML program where
+    the body is an `SSACFG` region. It must be terminated by a `return` op which
+    yields values with the same arity and types as the `FunctionType` results
+    of the containing `func`.
+
+    This op is a `Symbol` but does not introduce a new `SymbolTable`. As such,
+    it cannot represent nested symbols.
+
+    Example:
+
+    ```mlir
+    ml_program.func private @some_extern(i32) -> i32
+    ml_program.func @compute(%arg0 : i32) -> i32 {
+      ml_program.return %arg0 : i32
+    }
+    ```
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name,
+                       TypeAttrOf<FunctionType>:$function_type,
+                       OptionalAttr<StrAttr>:$sym_visibility);
+  let regions = (region AnyRegion:$body);
+
+  let extraClassDeclaration = [{
+    //===------------------------------------------------------------------===//
+    // CallableOpInterface
+    //===------------------------------------------------------------------===//
+
+    /// Returns the region on the current operation that is callable. This may
+    /// return null in the case of an external callable object, e.g. an external
+    /// function.
+    ::mlir::Region *getCallableRegion() {
+      return isExternal() ? nullptr : &getBody();
+    }
+
+    /// Returns the results types that the callable region produces when
+    /// executed.
+    ArrayRef<Type> getCallableResults() { return getFunctionType().getResults(); }
+
+    //===------------------------------------------------------------------===//
+    // FunctionOpInterface Methods
+    //===------------------------------------------------------------------===//
+
+    /// Returns the argument types of this function.
+    ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }
+
+    /// Returns the result types of this function.
+    ArrayRef<Type> getResultTypes() { return getFunctionType().getResults(); }
+
+    //===------------------------------------------------------------------===//
+    // RegionKindInterface Methods
+    //===------------------------------------------------------------------===//
+    static ::mlir::RegionKind getRegionKind(unsigned index) {
+      return ::mlir::RegionKind::SSACFG;
+    }
+
+    //===------------------------------------------------------------------===//
+    // SymbolOpInterface Methods
+    //===------------------------------------------------------------------===//
+
+    bool isDeclaration() { return isExternal(); }
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// SubgraphOp
+//===----------------------------------------------------------------------===//
+
+def MLProgram_SubgraphOp : MLProgram_Op<"subgraph", [
+    CallableOpInterface, FunctionOpInterface, HasOnlyGraphRegion,
+    IsolatedFromAbove, RegionKindInterface, SingleBlock, Symbol
+  ]> {
+  let summary = "An function containing a single `Graph` region";
+  let description = [{
+    This simple function container represents callables in an ML program where
+    the body is a `Graph` region containing a single block. It must be
+    terminated by an `output` op which yields values with the same arity and
+    types as the `FunctionType` results of the containing `subgraph`.
+
+    This op is a `Symbol` but does not introduce a new `SymbolTable`. As such,
+    it cannot represented nested symbols.
+
+    Example:
+
+    ```mlir
+    ml_program.subgraph private @some_extern(i32) -> i32
+    ml_program.subgraph @compute(%arg0 : i32) -> i32 {
+      ml_program.output %arg0 : i32
+    }
+    ```
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name,
+                       TypeAttrOf<FunctionType>:$function_type,
+                       OptionalAttr<StrAttr>:$sym_visibility);
+  let regions = (region AnyRegion:$body);
+
+  let extraClassDeclaration = [{
+    //===------------------------------------------------------------------===//
+    // CallableOpInterface
+    //===------------------------------------------------------------------===//
+
+    /// Returns the region on the current operation that is callable. This may
+    /// return null in the case of an external callable object, e.g. an external
+    /// function.
+    ::mlir::Region *getCallableRegion() { return isExternal() ? nullptr : &getBody(); }
+
+    /// Returns the results types that the callable region produces when
+    /// executed.
+    ArrayRef<Type> getCallableResults() { return getFunctionType().getResults(); }
+
+    //===------------------------------------------------------------------===//
+    // FunctionOpInterface Methods
+    //===------------------------------------------------------------------===//
+
+    /// Returns the argument types of this function.
+    ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }
+
+    /// Returns the result types of this function.
+    ArrayRef<Type> getResultTypes() { return getFunctionType().getResults(); }
+
+    //===------------------------------------------------------------------===//
+    // SymbolOpInterface Methods
+    //===------------------------------------------------------------------===//
+
+    bool isDeclaration() { return isExternal(); }
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// OutputOp
+//===----------------------------------------------------------------------===//
+
+def MLProgram_OutputOp : MLProgram_Op<"output", [
+    NoSideEffect, HasParent<"SubgraphOp">, ReturnLike, Terminator
+  ]> {
+  let summary = "Outputs values from a subgraph function";
+  let description = [{
+    The `output` operation terminates a subgraph by yielding values
+    to the caller.
+    The operation takes variable number of operands and produces no results.
+    The operand number and types must match the signature of the function
+    that contains the operation.
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$operands);
+
+  let builders = [OpBuilder<(ins), [{
+    build($_builder, $_state, llvm::None);
+  }]>];
+
+  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
+  let hasVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// ReturnOp
+//===----------------------------------------------------------------------===//
+
+def MLProgram_ReturnOp : MLProgram_Op<"return", [
+    NoSideEffect, HasParent<"FuncOp">, ReturnLike, Terminator
+  ]> {
+  let summary = "Returns values from a `func` function";
+  let description = [{
+    The `return` operation terminates a `func` function by yielding values
+    to the caller.
+    The operation takes variable number of operands and produces no results.
+    The operand number and types must match the signature of the function
+    that contains the operation.
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$operands);
+
+  let builders = [OpBuilder<(ins), [{
+    build($_builder, $_state, llvm::None);
+  }]>];
+
+  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
+  let hasVerifier = 1;
+}
+
+#endif // MLPROGRAM_OPS

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/EmitC/IR/EmitC.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.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"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/Quant/QuantOps.h"
 #include "mlir/Dialect/SCF/BufferizableOpInterfaceImpl.h"
registry.insert<acc::OpenACCDialect,
                   DLTIDialect,
                   emitc::EmitCDialect,
                   func::FuncDialect,
                   gpu::GPUDialect,
                   LLVM::LLVMDialect,
                   linalg::LinalgDialect,
                   math::MathDialect,
                   memref::MemRefDialect,
+                  ml_program::MLProgramDialect,
                   scf::SCFDialect,
                   omp::OpenMPDialect,
                   pdl::PDLDialect,
                   pdl_interp::PDLInterpDialect,
                   quant::QuantizationDialect,
                   spirv::SPIRVDialect,
                   arm_sve::ArmSVEDialect,
                   vector::VectorDialect,

mlir/lib/Dialect/CMakeLists.txt

 add_subdirectory(DLTI)
 add_subdirectory(EmitC)
 add_subdirectory(Func)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(Math)
 add_subdirectory(MemRef)
+add_subdirectory(MLProgram)
 add_subdirectory(OpenACC)
 add_subdirectory(OpenMP)
 add_subdirectory(PDL)
 add_subdirectory(PDLInterp)
 add_subdirectory(Quant)
 add_subdirectory(SCF)
 add_subdirectory(Shape)
 add_subdirectory(SparseTensor)

mlir/lib/Dialect/MLProgram/CMakeLists.txt

+add_subdirectory(IR)

mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt

+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
+  )

mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp

+//===- MLProgramDialect.cpp - 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"
+      >();
+}

mlir/lib/Dialect/MLProgram/IR/MLProgramOps.cpp

+//===- MLProgramOps.cpp - 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"
+#include "mlir/IR/FunctionImplementation.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"
+
+//===----------------------------------------------------------------------===//
+// FuncOp
+//===----------------------------------------------------------------------===//
+
+ParseResult FuncOp::parse(OpAsmParser &parser, OperationState &result) {
+  auto buildFuncType =
+      [](Builder &builder, ArrayRef<Type> argTypes, ArrayRef<Type> results,
+         function_interface_impl::VariadicFlag,
+         std::string &) { return builder.getFunctionType(argTypes, results); };
+
+  return function_interface_impl::parseFunctionOp(
+      parser, result, /*allowVariadic=*/false, buildFuncType);
+}
+
+void FuncOp::print(OpAsmPrinter &p) {
+  function_interface_impl::printFunctionOp(p, *this, /*isVariadic=*/false);
+}
+
+//===----------------------------------------------------------------------===//
+// SubgraphOp
+//===----------------------------------------------------------------------===//
+
+ParseResult SubgraphOp::parse(OpAsmParser &parser, OperationState &result) {
+  auto buildFuncType =
+      [](Builder &builder, ArrayRef<Type> argTypes, ArrayRef<Type> results,
+         function_interface_impl::VariadicFlag,
+         std::string &) { return builder.getFunctionType(argTypes, results); };
+
+  return function_interface_impl::parseFunctionOp(
+      parser, result, /*allowVariadic=*/false, buildFuncType);
+}
+
+void SubgraphOp::print(OpAsmPrinter &p) {
+  function_interface_impl::printFunctionOp(p, *this, /*isVariadic=*/false);
+}
+
+//===----------------------------------------------------------------------===//
+// OutputOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult OutputOp::verify() {
+  auto function = cast<SubgraphOp>((*this)->getParentOp());
+
+  // The operand number and types must match the function signature.
+  const auto &results = function.getFunctionType().getResults();
+  if (getNumOperands() != results.size())
+    return emitOpError("has ")
+           << getNumOperands() << " operands, but enclosing function (@"
+           << function.getName() << ") outputs " << results.size();
+
+  for (unsigned i = 0, e = results.size(); i != e; ++i)
+    if (getOperand(i).getType() != results[i])
+      return emitError() << "type of output operand " << i << " ("
+                         << getOperand(i).getType()
+                         << ") doesn't match function result type ("
+                         << results[i] << ")"
+                         << " in function @" << function.getName();
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// ReturnOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult ReturnOp::verify() {
+  auto function = cast<FuncOp>((*this)->getParentOp());
+
+  // The operand number and types must match the function signature.
+  const auto &results = function.getFunctionType().getResults();
+  if (getNumOperands() != results.size())
+    return emitOpError("has ")
+           << getNumOperands() << " operands, but enclosing function (@"
+           << function.getName() << ") returns " << results.size();
+
+  for (unsigned i = 0, e = results.size(); i != e; ++i)
+    if (getOperand(i).getType() != results[i])
+      return emitError() << "type of return operand " << i << " ("
+                         << getOperand(i).getType()
+                         << ") doesn't match function result type ("
+                         << results[i] << ")"
+                         << " in function @" << function.getName();
+
+  return success();
+}

mlir/test/Dialect/MLProgram/invalid.mlir

+// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -verify-diagnostics %s
+
+ml_program.func @ssa_enforced(%arg0 : i32) -> i32 {
+  // expected-error @+1 {{does not dominate this use}}
+  %1 = "unregistered.dummy"(%0) : (i32) -> i32
+  // expected-note @+1 {{operand defined here}}
+  %0 = "unregistered.dummy"(%arg0) : (i32) -> i32
+  ml_program.return %0 : i32
+}
+
+// -----
+ml_program.func @return_arity_match(%arg0 : i32) -> i32 {
+  // expected-error @+1 {{enclosing function (@return_arity_match) returns 1}}
+  ml_program.return %arg0, %arg0 : i32, i32
+}
+
+// -----
+ml_program.func @return_type_match(%arg0 : i64) -> i32 {
+  // expected-error @+1 {{doesn't match function result}}
+  ml_program.return %arg0 : i64
+}
+
+// -----
+ml_program.subgraph @output_arity_match(%arg0 : i32) -> i32 {
+  // expected-error @+1 {{enclosing function (@output_arity_match) outputs 1}}
+  ml_program.output %arg0, %arg0 : i32, i32
+}
+
+// -----
+ml_program.subgraph @output_type_match(%arg0 : i64) -> i32 {
+  // expected-error @+1 {{doesn't match function result}}
+  ml_program.output %arg0 : i64
+}

mlir/test/Dialect/MLProgram/ops.mlir

+// RUN: mlir-opt %s --allow-unregistered-dialect | mlir-opt --allow-unregistered-dialect | FileCheck %s
+// RUN: mlir-opt %s --allow-unregistered-dialect --mlir-print-op-generic | mlir-opt --allow-unregistered-dialect | FileCheck %s
+
+// CHECK-LABEL: ml_program.func private @extern_func
+ml_program.func private @extern_func(i32) -> i32
+
+// CHECK-LABEL: ml_program.func @defined_func
+ml_program.func @defined_func(%arg0 : i32) -> i32 {
+  ml_program.return %arg0 : i32
+}
+
+// CHECK-LABEL: ml_program.subgraph private @extern_subgraph
+ml_program.subgraph private @extern_subgraph(i32) -> i32
+
+// CHECK-LABEL: ml_program.subgraph @compute_subgraph
+ml_program.subgraph @compute_subgraph(%arg0 : i32) -> i32 {
+  %1 = "unregistered.dummy"(%0) : (i32) -> i32
+  %0 = "unregistered.dummy"(%arg0) : (i32) -> i32
+  ml_program.output %0 : i32
+}

mlir/test/mlir-opt/commandline.mlir

 // CHECK-NEXT: dlti
 // CHECK-NEXT: emitc
 // CHECK-NEXT: func
 // CHECK-NEXT: gpu
 // CHECK-NEXT: linalg
 // CHECK-NEXT: llvm
 // CHECK-NEXT: math
 // CHECK-NEXT: memref
+// CHECK-NEXT: ml_program
 // CHECK-NEXT: nvvm
 // CHECK-NEXT: omp
 // CHECK-NEXT: pdl
 // CHECK-NEXT: pdl_interp
 // CHECK-NEXT: quant
 // CHECK-NEXT: rocdl
 // CHECK-NEXT: scf
 // CHECK-NEXT: shape
 // CHECK-NEXT: sparse_tensor
 // CHECK-NEXT: spv
 // CHECK-NEXT: tensor
 // CHECK-NEXT: test
 // CHECK-NEXT: tosa
 // CHECK-NEXT: vector
 // CHECK-NEXT: x86vector

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":LLVMIRTransforms",
         ":LLVMPassIncGen",
         ":LinalgOps",
         ":LinalgPassIncGen",
         ":LinalgToLLVM",
         ":LinalgToSPIRV",
         ":LinalgToStandard",
         ":LinalgTransforms",
+        ":MLProgramDialect",
         ":MathDialect",
         ":MathToLLVM",
         ":MathToLibm",
         ":MathToSPIRV",
         ":MathTransforms",
         ":MemRefDialect",
         ":MemRefToLLVM",
         ":MemRefToSPIRV",
deps = [
         ":Pass",
         ":TensorDialect",
         ":Transforms",
         ":VectorOps",
         "//llvm:Support",
     ],
 )
 
+##---------------------------------------------------------------------------##
+# MLProgram dialect
+##---------------------------------------------------------------------------##
+
+td_library(
+    name = "MLProgramOpsTdFiles",
+    srcs = [
+        "include/mlir/Dialect/MLProgram/IR/MLProgramBase.td",
+        "include/mlir/Dialect/MLProgram/IR/MLProgramOps.td",
+    ],
+    includes = ["include"],
+    deps = [
+        ":CallInterfacesTdFiles",
+        ":ControlFlowInterfacesTdFiles",
+        ":FunctionInterfacesTdFiles",
+        ":OpBaseTdFiles",
+        ":RegionKindInterfaceIncGen",
+        ":SideEffectInterfacesTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "MLProgramOpsIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-op-decls"],
+            "include/mlir/Dialect/MLProgram/IR/MLProgramOps.h.inc",
+        ),
+        (
+            ["-gen-op-defs"],
+            "include/mlir/Dialect/MLProgram/IR/MLProgramOps.cpp.inc",
+        ),
+        (
+            ["-gen-dialect-decls"],
+            "include/mlir/Dialect/MLProgram/IR/MLProgramOpsDialect.h.inc",
+        ),
+        (
+            ["-gen-dialect-defs"],
+            "include/mlir/Dialect/MLProgram/IR/MLProgramOpsDialect.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/MLProgram/IR/MLProgramOps.td",
+    deps = [":MLProgramOpsTdFiles"],
+)
+
+cc_library(
+    name = "MLProgramDialect",
+    srcs = glob([
+        "lib/Dialect/MLProgram/IR/*.cpp",
+        "lib/Dialect/MLProgram/IR/*.h",
+    ]),
+    hdrs = glob([
+        "include/mlir/Dialect/MLProgram/IR/*.h",
+    ]),
+    includes = ["include"],
+    deps = [
+        ":ControlFlowInterfaces",
+        ":IR",
+        ":MLProgramOpsIncGen",
+        ":Pass",
+        ":Support",
+        "//llvm:Support",
+    ],
+)
+
+##---------------------------------------------------------------------------##
+# Allocation interfaces
+##---------------------------------------------------------------------------##
+
 td_library(
     name = "AllocationOpInterfaceTdFiles",
     srcs = ["include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.td"],
     includes = ["include"],
     deps = [":OpBaseTdFiles"],
 )
 
 gentbl_cc_library(