1//===- MLProgramOps.td - Structural ML Program Ops ---------*- tablegen -*-===// 2// 3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4// See https://llvm.org/LICENSE.txt for license information. 5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6// 7//===----------------------------------------------------------------------===// 8 9#ifndef MLPROGRAM_OPS 10#define MLPROGRAM_OPS 11 12include "mlir/Dialect/MLProgram/IR/MLProgramBase.td" 13include "mlir/Dialect/MLProgram/IR/MLProgramAttributes.td" 14include "mlir/Dialect/MLProgram/IR/MLProgramTypes.td" 15include "mlir/Interfaces/CallInterfaces.td" 16include "mlir/Interfaces/ControlFlowInterfaces.td" 17include "mlir/Interfaces/SideEffectInterfaces.td" 18include "mlir/Interfaces/FunctionInterfaces.td" 19include "mlir/IR/OpAsmInterface.td" 20include "mlir/IR/RegionKindInterface.td" 21include "mlir/IR/SymbolInterfaces.td" 22 23class MLProgram_Op<string mnemonic, list<Trait> traits = []> : 24 Op<MLProgram_Dialect, mnemonic, traits>; 25 26//===----------------------------------------------------------------------===// 27// FuncOp 28//===----------------------------------------------------------------------===// 29 30def MLProgram_FuncOp : MLProgram_Op<"func", [ 31 FunctionOpInterface, IsolatedFromAbove, 32 RegionKindInterface, Symbol 33 ]> { 34 let summary = "Function containing a single `SSACFG` region"; 35 let description = [{ 36 This simple function container represents callables in an ML program where 37 the body is an `SSACFG` region. It must be terminated by a `return` op which 38 yields values with the same arity and types as the `FunctionType` results 39 of the containing `func`. 40 41 This op is a `Symbol` but does not introduce a new `SymbolTable`. As such, 42 it cannot represent nested symbols. 43 44 Example: 45 46 ```mlir 47 ml_program.func private @some_extern(i32) -> i32 48 ml_program.func @compute(%arg0 : i32) -> i32 { 49 ml_program.return %arg0 : i32 50 } 51 ``` 52 }]; 53 54 let arguments = (ins SymbolNameAttr:$sym_name, 55 TypeAttrOf<FunctionType>:$function_type, 56 OptionalAttr<DictArrayAttr>:$arg_attrs, 57 OptionalAttr<DictArrayAttr>:$res_attrs, 58 OptionalAttr<StrAttr>:$sym_visibility); 59 let regions = (region AnyRegion:$body); 60 61 let extraClassDeclaration = [{ 62 //===------------------------------------------------------------------===// 63 // FunctionOpInterface Methods 64 //===------------------------------------------------------------------===// 65 66 /// Returns the region on the current operation that is callable. This may 67 /// return null in the case of an external callable object, e.g. an external 68 /// function. 69 ::mlir::Region *getCallableRegion() { 70 return isExternal() ? nullptr : &getBody(); 71 } 72 73 /// Returns the argument types of this function. 74 ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); } 75 76 /// Returns the result types of this function. 77 ArrayRef<Type> getResultTypes() { return getFunctionType().getResults(); } 78 79 //===------------------------------------------------------------------===// 80 // RegionKindInterface Methods 81 //===------------------------------------------------------------------===// 82 static ::mlir::RegionKind getRegionKind(unsigned index) { 83 return ::mlir::RegionKind::SSACFG; 84 } 85 86 //===------------------------------------------------------------------===// 87 // SymbolOpInterface Methods 88 //===------------------------------------------------------------------===// 89 90 bool isDeclaration() { return isExternal(); } 91 }]; 92 93 let hasCustomAssemblyFormat = 1; 94} 95 96//===----------------------------------------------------------------------===// 97// GlobalOp 98//===----------------------------------------------------------------------===// 99 100def MLProgram_GlobalOp : MLProgram_Op<"global", [ 101 Symbol 102 ]> { 103 let summary = "Module level declaration of a global variable"; 104 let description = [{ 105 Declares a named global variable (or constant). 106 107 A global contains a value of a specified type which can be accessed at 108 runtime via appropriate load/store operations. It can be mutable or 109 constant, optionally taking an initial value or declared as 110 extern (in which case, the initial value is found in external storage 111 by symbol name). 112 113 Generally, the type of the global and the type of the initial value 114 will be the same. However, for type hierarchies which can have a more 115 generalized bounding type that can be assigned from a narrow type, this 116 is allowed (but not verified). 117 118 Examples: 119 120 ```mlir 121 // Constant global. 122 ml_program.global @foobar(dense<4> : tensor<4xi32>) : tensor<?xi32> 123 124 // Constant with external linkage. 125 ml_program.global mutable @foobar(#ml_program.extern<tensor<4xi32>>) 126 : tensor<?xi32> 127 128 // Mutable global with an undefined initial value. 129 ml_program.global mutable @foobar : tensor<?xi32> 130 ``` 131 }]; 132 133 let arguments = (ins 134 SymbolNameAttr:$sym_name, 135 TypeAttr:$type, 136 UnitAttr:$is_mutable, 137 OptionalAttr<AnyAttr>:$value, 138 OptionalAttr<StrAttr>:$sym_visibility 139 ); 140 141 let assemblyFormat = [{ 142 custom<SymbolVisibility>($sym_visibility) 143 (`mutable` $is_mutable^)? 144 $sym_name `` 145 custom<TypedInitialValue>($type, $value) 146 attr-dict 147 }]; 148 149 let hasVerifier = 1; 150} 151 152//===----------------------------------------------------------------------===// 153// GlobalLoadOp 154//===----------------------------------------------------------------------===// 155 156def MLProgram_GlobalLoadOp : MLProgram_Op<"global_load", [ 157 DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>, 158 DeclareOpInterfaceMethods<SymbolUserOpInterface> 159 ]> { 160 let summary = "Direct load of a mutable value from a global"; 161 let description = [{ 162 Performs a non-atomic, non-volatile, non-synchronized load from a global 163 that may be mutable. 164 165 It is fully expected that these constraints are not suitable for 166 all situations, and alternative ops should be defined and used for more 167 advanced cases. 168 169 This op is side effecting and may not be valid to use in graph regions 170 without additional consideration to evaluation order constraints. See 171 `global_load_graph` for op which allows for explicit ordering constraints. 172 173 Example: 174 175 ```mlir 176 %0 = ml_program.global_load @foobar : tensor<?xi32> 177 ``` 178 }]; 179 180 let arguments = (ins 181 Arg<SymbolRefAttr, "", [MemRead]>:$global 182 ); 183 let results = (outs 184 AnyType:$result 185 ); 186 187 let assemblyFormat = [{ 188 $global `:` type($result) attr-dict 189 }]; 190 191 let extraClassDeclaration = [{ 192 /// Gets the corresponding GlobalOp (or nullptr). 193 GlobalOp getGlobalOp(SymbolTableCollection &symbolTable); 194 }]; 195 196 let extraClassDefinition = [{ 197 void $cppClass::getAsmResultNames( 198 function_ref<void(::mlir::Value, ::llvm::StringRef)> setNameFn) { 199 setNameFn(getResult(), getGlobal().getLeafReference()); 200 } 201 }]; 202} 203 204//===----------------------------------------------------------------------===// 205// GlobalLoadConstOp 206//===----------------------------------------------------------------------===// 207 208def MLProgram_GlobalLoadConstOp : MLProgram_Op<"global_load_const", [ 209 Pure, 210 DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>, 211 DeclareOpInterfaceMethods<SymbolUserOpInterface> 212 ]> { 213 let summary = "Direct load a constant value from a global"; 214 let description = [{ 215 Loads a constant (immutable) value from a global directly by symbol. 216 217 This op is only legal for globals that are not mutable and exists because 218 such a load can be considered to have no side effects. 219 220 Example: 221 222 ```mlir 223 %0 = ml_program.global_load_const @foobar : tensor<?xi32> 224 ``` 225 }]; 226 227 let arguments = (ins 228 SymbolRefAttr:$global 229 ); 230 let results = (outs 231 AnyType:$result 232 ); 233 234 let assemblyFormat = [{ 235 $global `:` type($result) attr-dict 236 }]; 237 238 let extraClassDeclaration = [{ 239 /// Gets the corresponding GlobalOp (or nullptr). 240 GlobalOp getGlobalOp(SymbolTableCollection &symbolTable); 241 }]; 242 243 let extraClassDefinition = [{ 244 void $cppClass::getAsmResultNames( 245 function_ref<void(::mlir::Value, ::llvm::StringRef)> setNameFn) { 246 setNameFn(getResult(), getGlobal().getLeafReference()); 247 } 248 }]; 249} 250 251//===----------------------------------------------------------------------===// 252// GlobalLoadGraphOp 253//===----------------------------------------------------------------------===// 254 255def MLProgram_GlobalLoadGraphOp : MLProgram_Op<"global_load_graph", [ 256 DeclareOpInterfaceMethods<SymbolUserOpInterface> 257 ]> { 258 let summary = "Direct load of a mutable value from a global in Graph region"; 259 let description = [{ 260 Performs a non-atomic, non-volatile, non-synchronized load from a global 261 that may be mutable. 262 263 It is fully expected that these constraints are not suitable for all 264 situations, and alternative ops should be defined and used for more advanced 265 cases. 266 267 This op is side effecting and may not be valid to use in graph regions 268 without additional consideration to evaluation order constraints. 269 270 Example: 271 272 ```mlir 273 %0, %cstr = ml_program.global_load_graph @foobar 274 ordering (%token -> !ml_program.token) : tensor<?xi32> 275 ``` 276 }]; 277 278 let arguments = (ins 279 Arg<SymbolRefAttr, "", [MemRead]>:$global, 280 Variadic<MLProgram_TokenType>:$consumeTokens 281 ); 282 let results = (outs 283 AnyType:$result, 284 MLProgram_TokenType:$produceToken 285 ); 286 287 let assemblyFormat = [{ 288 $global `` custom<TokenOrdering>($consumeTokens, type($produceToken)) `:` type($result) attr-dict 289 }]; 290 291 let extraClassDeclaration = [{ 292 /// Gets the corresponding GlobalOp (or nullptr). 293 GlobalOp getGlobalOp(SymbolTableCollection &symbolTable); 294 }]; 295} 296 297//===----------------------------------------------------------------------===// 298// GlobalStoreOp 299//===----------------------------------------------------------------------===// 300 301def MLProgram_GlobalStoreOp : MLProgram_Op<"global_store", [ 302 DeclareOpInterfaceMethods<SymbolUserOpInterface> 303 ]> { 304 let summary = "Direct store of a value into a mutable global"; 305 let description = [{ 306 Performs a non-atomic, non-volatile, non-synchronized store to a mutable 307 global. 308 309 It is fully expected that these constraints are not suitable for 310 all situations, and alternative ops should be defined and used for more 311 advanced cases. 312 313 This op is side effecting and may not be valid to use in graph regions 314 without additional consideration to evaluation order constraints. See 315 `global_store_graph` for op which allows for explicit ordering constraints. 316 317 Example: 318 319 ```mlir 320 ml_program.global_store @foobar = %0 : tensor<?xi32> 321 ``` 322 }]; 323 324 let arguments = (ins 325 Arg<SymbolRefAttr, "", [MemWrite]>:$global, 326 AnyType:$value 327 ); 328 329 let assemblyFormat = [{ 330 $global `=` $value `:` type($value) attr-dict 331 }]; 332 333 let extraClassDeclaration = [{ 334 /// Gets the corresponding GlobalOp (or nullptr). 335 GlobalOp getGlobalOp(SymbolTableCollection &symbolTable); 336 }]; 337} 338 339//===----------------------------------------------------------------------===// 340// GlobalStoreGraphOp 341//===----------------------------------------------------------------------===// 342 343def MLProgram_GlobalStoreGraphOp : MLProgram_Op<"global_store_graph", [ 344 DeclareOpInterfaceMethods<SymbolUserOpInterface> 345 ]> { 346 let summary = "Direct store of a value into a mutable global"; 347 let description = [{ 348 Performs a non-atomic, non-volatile, non-synchronized store to a mutable 349 global. 350 351 It is fully expected that these constraints are not suitable for 352 all situations, and alternative ops should be defined and used for more 353 advanced cases. 354 355 This op is side effecting and may not be valid to use in graph regions 356 without additional consideration to evaluation order constraints. 357 358 Example: 359 360 ```mlir 361 %token = ml_program.global_store @foobar = %0 : tensor<?xi32> 362 ordering (%in_token -> !ml_program.token) : tensor<?xi32> 363 ``` 364 }]; 365 366 let arguments = (ins 367 Arg<SymbolRefAttr, "", [MemRead]>:$global, 368 AnyType:$value, 369 Variadic<MLProgram_TokenType>:$consumeTokens 370 ); 371 let results = (outs 372 MLProgram_TokenType:$produceToken 373 ); 374 375 let assemblyFormat = [{ 376 $global `=` $value `` custom<TokenOrdering>($consumeTokens, type($produceToken)) `:` type($value) attr-dict 377 }]; 378 379 let extraClassDeclaration = [{ 380 /// Gets the corresponding GlobalOp (or nullptr). 381 GlobalOp getGlobalOp(SymbolTableCollection &symbolTable); 382 }]; 383} 384 385//===----------------------------------------------------------------------===// 386// SubgraphOp 387//===----------------------------------------------------------------------===// 388 389def MLProgram_SubgraphOp : MLProgram_Op<"subgraph", [ 390 FunctionOpInterface, HasOnlyGraphRegion, 391 IsolatedFromAbove, RegionKindInterface, SingleBlock, Symbol 392 ]> { 393 let summary = "An function containing a single `Graph` region"; 394 let description = [{ 395 This simple function container represents callables in an ML program where 396 the body is a `Graph` region containing a single block. It must be 397 terminated by an `output` op which yields values with the same arity and 398 types as the `FunctionType` results of the containing `subgraph`. 399 400 This op is a `Symbol` but does not introduce a new `SymbolTable`. As such, 401 it cannot represented nested symbols. 402 403 Example: 404 405 ```mlir 406 ml_program.subgraph private @some_extern(i32) -> i32 407 ml_program.subgraph @compute(%arg0 : i32) -> i32 { 408 ml_program.output %arg0 : i32 409 } 410 ``` 411 }]; 412 413 let arguments = (ins SymbolNameAttr:$sym_name, 414 TypeAttrOf<FunctionType>:$function_type, 415 OptionalAttr<DictArrayAttr>:$arg_attrs, 416 OptionalAttr<DictArrayAttr>:$res_attrs, 417 OptionalAttr<StrAttr>:$sym_visibility); 418 let regions = (region AnyRegion:$body); 419 420 let extraClassDeclaration = [{ 421 //===------------------------------------------------------------------===// 422 // FunctionOpInterface Methods 423 //===------------------------------------------------------------------===// 424 425 /// Returns the region on the current operation that is callable. This may 426 /// return null in the case of an external callable object, e.g. an external 427 /// function. 428 ::mlir::Region *getCallableRegion() { return isExternal() ? nullptr : &getBody(); } 429 430 /// Returns the argument types of this function. 431 ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); } 432 433 /// Returns the result types of this function. 434 ArrayRef<Type> getResultTypes() { return getFunctionType().getResults(); } 435 436 //===------------------------------------------------------------------===// 437 // SymbolOpInterface Methods 438 //===------------------------------------------------------------------===// 439 440 bool isDeclaration() { return isExternal(); } 441 }]; 442 443 let hasCustomAssemblyFormat = 1; 444} 445 446//===----------------------------------------------------------------------===// 447// OutputOp 448//===----------------------------------------------------------------------===// 449 450def MLProgram_OutputOp : MLProgram_Op<"output", [ 451 Pure, HasParent<"SubgraphOp">, ReturnLike, Terminator 452 ]> { 453 let summary = "Outputs values from a subgraph function"; 454 let description = [{ 455 The `output` operation terminates a subgraph by yielding values 456 to the caller. 457 The operation takes variable number of operands and produces no results. 458 The operand number and types must match the signature of the function 459 that contains the operation. 460 }]; 461 462 let arguments = (ins Variadic<AnyType>:$operands); 463 464 let builders = [OpBuilder<(ins), [{ 465 build($_builder, $_state, std::nullopt); 466 }]>]; 467 468 let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?"; 469 let hasVerifier = 1; 470} 471 472//===----------------------------------------------------------------------===// 473// ReturnOp 474//===----------------------------------------------------------------------===// 475 476def MLProgram_ReturnOp : MLProgram_Op<"return", [ 477 Pure, HasParent<"FuncOp">, ReturnLike, Terminator 478 ]> { 479 let summary = "Returns values from a `func` function"; 480 let description = [{ 481 The `return` operation terminates a `func` function by yielding values 482 to the caller. 483 The operation takes variable number of operands and produces no results. 484 The operand number and types must match the signature of the function 485 that contains the operation. 486 }]; 487 488 let arguments = (ins Variadic<AnyType>:$operands); 489 490 let builders = [OpBuilder<(ins), [{ 491 build($_builder, $_state, std::nullopt); 492 }]>]; 493 494 let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?"; 495 let hasVerifier = 1; 496} 497 498//===----------------------------------------------------------------------===// 499// TokenOp 500//===----------------------------------------------------------------------===// 501 502def MLProgram_TokenOp : MLProgram_Op<"token", [ 503 Pure 504 ]> { 505 let summary = "Produces a new token value"; 506 let description = [{ 507 Token values are used to chain side effecting ops in a graph so as to 508 establish an execution order. This op produces a token. 509 }]; 510 511 let results = (outs 512 MLProgram_TokenType:$token 513 ); 514 515 let assemblyFormat = "attr-dict"; 516} 517 518#endif // MLPROGRAM_OPS 519