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