Iterators: Create iterator state type and extract/insert/undef ops. (#606)

Author: ingomueller-net

experimental/iterators/include/iterators/Dialect/Iterators/IR/IteratorsOps.td

@@ -303,4 +303,94 @@ def Iterators_SinkOp : Iterators_Base_Op<"sink"> {
   let arguments = (ins Iterators_StreamOfLLVMStructOfNumerics:$input);
 }
 
+//===----------------------------------------------------------------------===//
+// Ops related to Iterator bodies.
+//===----------------------------------------------------------------------===//
+
+def Iterators_UndefStateOp : Iterators_Base_Op<"undefstate", [NoSideEffect]> {
+  let summary = "Create an undefined iterator state";
+  let results = (outs Iterators_State:$result);
+  let assemblyFormat = "attr-dict `:` qualified(type($result))";
+  let description = [{
+    Creates an iterator state of the given type with undefined field values.
+    All fields have to be set individually with `insertvalue` before the whole
+    state is fully defined.
+
+    This is similar to `llvm.undef` for `llvm.struct`.
+
+    Example:
+
+    ```
+    %undef_state = iterators.undefstate : !iterators.state<i32, tensor<?xi32>>
+    ```
+  }];
+}
+
+def Iterators_ExtractValueOp : Iterators_Base_Op<"extractvalue", [
+    NoSideEffect,
+    PredOpTrait<"index must exist in state",
+      CPred<"static_cast<uint64_t>($index.cast<IntegerAttr>().getInt())"
+            "    < $state.getType().cast<StateType>().getFieldTypes().size()">>,
+    AllMatch<["$state.getType().cast<StateType>().getFieldTypes()"
+              "    [$index.cast<IntegerAttr>().getInt()]",
+              "$result.getType()"],
+             "the return type must match the field type at the given index">,
+    DeclareOpInterfaceMethods<InferTypeOpInterface>
+    ]> {
+  let summary = "Extract the field value of the state";
+  let arguments = (ins Iterators_State:$state, IndexAttr:$index);
+  let results = (outs AnyType:$result);
+  let assemblyFormat = [{
+    $state `[` $index `]` attr-dict `:` qualified(type($state))
+  }];
+  let description = [{
+    Extracts the value of the given iterator state at the given index.
+
+    This is similar to `llvm.extractvalue` for `llvm.struct`.
+
+    Example:
+
+    ```
+    %state = ...
+    %value = iterators.extractvalue %state[0] :
+                 !iterators.state<i32, tensor<?xi32>>
+    ```
+  }];
+}
+
+def Iterators_InsertValueOp : Iterators_Base_Op<"insertvalue", [
+    NoSideEffect,
+    PredOpTrait<
+      "index must exist in state",
+      CPred<"static_cast<uint64_t>($index.cast<IntegerAttr>().getInt())"
+            "    < $state.getType().cast<StateType>().getFieldTypes().size()">>,
+    AllMatch<["$state.getType().cast<StateType>().getFieldTypes()"
+              "    [$index.cast<IntegerAttr>().getInt()]",
+              "$value.getType()"],
+             "the value type must match the field type at the given index">,
+    AllTypesMatch<["state", "result"]>
+    ]> {
+  let summary = "Insert a field value into the state";
+  let arguments = (ins Iterators_State:$state, IndexAttr:$index, AnyType:$value);
+  let results = (outs Iterators_State:$result);
+  let assemblyFormat = [{
+    $value `into` $state `[` $index `]` attr-dict `:` qualified(type($state))
+      custom<InsertValueType>(type($value), ref(type($state)), ref($index))
+  }];
+  let description = [{
+    Inserts the given value into the given iterator state at the given index.
+
+    This is similar to `llvm.insertvalue` for `llvm.struct`.
+
+    Example:
+
+    ```
+    %state = ...
+    %value = ...
+    %updated_state = iterators.insertvalue %state[0] (%value : i32) :
+                         !iterators.state<i32, tensor<?xi32>>
+    ```
+  }];
+}
+
 #endif // ITERATORS_DIALECT_ITERATORS_IR_ITERATORSOPS

experimental/iterators/include/iterators/Dialect/Iterators/IR/IteratorsTypes.td

@@ -154,4 +154,32 @@ class Iterators_IsStreamOf<string name, Type type>
 class Iterators_IsStreamOfLLVMStructOfNumericsPred<string name>
   : Iterators_IsStreamOf<name, Iterators_StreamOfLLVMStructOfNumerics>;
 
+//===----------------------------------------------------------------------===//
+// Types related to Iterator bodies.
+//===----------------------------------------------------------------------===//
+
+def Iterators_State : Iterators_Type<"State", "state"> {
+  let summary = "State of an iterator used by its body";
+  let parameters =  (ins ArrayRefParameter<"Type", "list of types">:$fieldTypes);
+  let assemblyFormat = "`<` qualified($fieldTypes) `>`";
+  let description = [{
+    An iterator state is a collection of values identified by ordinal numbers,
+    i.e., an (unnamed but typed) tuple. The values are referred to as "fields";
+    their types are referred to as "field types". An iterator state is used by
+    iterator bodies, i.e., by the open, next, and close functions that implement
+    the logic that iterator ops get lowered to, and holds the state that is
+    required during the iteration (which gets passed around different calls to
+    open, next, and close).
+
+    This is similar to (anonymous) `llvm.struct` but allows for storing values
+    of arbitrary types.
+
+    Example:
+
+    ```
+    %undef_state = iterators.undefstate : !iterators.state<i32, tensor<?xi32>>
+    ```
+  }];
+}
+
 #endif // ITERATORS_DIALECT_ITERATORS_IR_ITERATORSTYPES

experimental/iterators/lib/Dialects/Iterators/IR/Iterators.cpp

@@ -45,9 +45,33 @@ void IteratorsDialect::initialize() {
 // Iterators operations
 //===----------------------------------------------------------------------===//
 
+static ParseResult parseInsertValueType(AsmParser & /*parser*/, Type &valueType,
+                                        Type stateType, IntegerAttr indexAttr) {
+  int64_t index = indexAttr.getValue().getSExtValue();
+  auto castedStateType = stateType.cast<StateType>();
+  valueType = castedStateType.getFieldTypes()[index];
+  return success();
+}
+
+static void printInsertValueType(AsmPrinter & /*printer*/, Operation * /*op*/,
+                                 Type /*valueType*/, Type /*stateType*/,
+                                 IntegerAttr /*indexAttr*/) {}
+
 #define GET_OP_CLASSES
 #include "iterators/Dialect/Iterators/IR/IteratorsOps.cpp.inc"
 
+LogicalResult ExtractValueOp::inferReturnTypes(
+    MLIRContext * /*context*/, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  auto stateType = operands[0].getType().cast<StateType>();
+  auto indexAttr = attributes.getAs<IntegerAttr>("index");
+  int64_t index = indexAttr.getValue().getSExtValue();
+  Type fieldType = stateType.getFieldTypes()[index];
+  inferredReturnTypes.assign({fieldType});
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Iterators types
 //===----------------------------------------------------------------------===//

experimental/iterators/test/Dialect/Iterators/state.mlir

@@ -0,0 +1,19 @@
+// Test that we can parse and verify ops on iterator state correctly, and that
+// they round-trip through assembly.
+// RUN: mlir-proto-opt %s \
+// RUN: | FileCheck %s
+
+func.func @testUndefInsertExtract() {
+// CHECK-LABEL: func.func @testUndefInsertExtract() {
+  %initial_state = iterators.undefstate : !iterators.state<i32>
+// CHECK-NEXT:     %[[V0:.*]] = iterators.undefstate : !iterators.state<i32>
+  %value = arith.constant 0 : i32
+// CHECK-NEXT:     %[[V1:.*]] = arith.constant 0 : i32
+  %inserted_state = iterators.insertvalue %value into %initial_state[0] : !iterators.state<i32>
+// CHECK-NEXT:     %[[V2:.*]] = iterators.insertvalue %[[V1]] into %[[V0]][0] : !iterators.state<i32>
+  %extracted_value = iterators.extractvalue %inserted_state[0] : !iterators.state<i32>
+// CHECK-NEXT:     %[[V3:.*]] = iterators.extractvalue %[[V2]][0] : !iterators.state<i32>
+  return
+// CHECK-NEXT:     return
+}
+// CHECK-NEXT:   }