Rebased to main

Author: PetrifiedPanda

docs/CoreDSLDialect.md

@@ -849,6 +849,72 @@ Traits: `HasParent<InstructionOp>`, `NoRegionArguments`, `SingleBlock`, `Termina
 
 
 
+### `coredsl.switch` (coredsl::SwitchOp)
+
+_Switch-case operation on arbitrary integers_
+
+Syntax:
+
+```
+operation ::= `coredsl.switch` $arg attr-dict `:` type($arg) (`->` type($results)^)?
+              custom<SwitchCases>($cases, $caseRegions) `\n`
+              `` `default` $defaultRegion
+```
+
+`coredsl.switch` is a modified version of `scf.index_switch` that branches
+to one of the given regions based on the values of the argument and the
+cases. As opposed to `scf.index_switch`, this operation accepts any integer
+as argument.
+
+The operation always has a "default" region and any number of case regions
+denoted by integer constants. Control-flow transfers to the case region
+whose constant value equals the value of the argument. If the argument does
+not equal any of the case values, control-flow transfer to the "default"
+region.
+
+Example:
+
+```mlir
+%0 = coredsl.switch %arg0 : ui32 -> ui32
+case 2 {
+  %1 = hwarith.constant 10 : ui32
+  coredsl.yield %1 : ui32
+}
+case 5 {
+  %2 = hwarith.constant 20 : ui32
+  coredsl.yield %2 : ui32
+}
+default {
+  %3 = hwarith.constant 30 : ui32
+  coredsl.yield %3 : ui32
+}
+```
+
+Traits: `RecursiveMemoryEffects`, `SingleBlockImplicitTerminator<coredsl::YieldOp>`, `SingleBlock`
+
+Interfaces: `RegionBranchOpInterface`
+
+#### Attributes:
+
+<table>
+<tr><th>Attribute</th><th>MLIR Type</th><th>Description</th></tr>
+<tr><td><code>cases</code></td><td>::mlir::ArrayAttr</td><td>array attribute</td></tr>
+</table>
+
+#### Operands:
+
+| Operand | Description |
+| :-----: | ----------- |
+| `arg` | integer |
+
+#### Results:
+
+| Result | Description |
+| :----: | ----------- |
+| `results` | variadic of any type |
+
+
+
 ### `coredsl.xor` (coredsl::XorOp)
 
 _Bitwise XOR operator._
@@ -890,6 +956,39 @@ Effects: `MemoryEffects::Effect{}`
 
 
 
+### `coredsl.yield` (coredsl::YieldOp)
+
+_Loop yield and termination operation_
+
+Syntax:
+
+```
+operation ::= `coredsl.yield` attr-dict ($results^ `:` type($results))?
+```
+
+The `coredsl.yield` operation is equivalent to the `scf.yield` operation,
+but only works for `coredsl.switch`
+If `coredsl.yield` has any operands, the operands must match the parent
+operation's results.
+If the parent operation defines no values, then the `coredsl.yield` may be
+left out in the custom syntax and the builders will insert one implicitly.
+Otherwise, it has to be present in the syntax to indicate which values are
+yielded.
+
+Traits: `AlwaysSpeculatableImplTrait`, `HasParent<SwitchOp>`, `ReturnLike`, `Terminator`
+
+Interfaces: `ConditionallySpeculatable`, `NoMemoryEffect (MemoryEffectOpInterface)`, `RegionBranchTerminatorOpInterface`
+
+Effects: `MemoryEffects::Effect{}`
+
+#### Operands:
+
+| Operand | Description |
+| :-----: | ----------- |
+| `results` | variadic of any type |
+
+
+
 ## Enums
 
 ### AddressSpaceAccessMode

include/shortnail/Dialect/CoreDSL/CoreDSLOps.td

@@ -685,4 +685,92 @@ def CastOp : CoreDSL_Op<"cast", [Pure]> {
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// YieldOp
+//===----------------------------------------------------------------------===//
+
+def YieldOp : CoreDSL_Op<"yield", [Pure, ReturnLike, Terminator,
+    ParentOneOf<["SwitchOp"]>]> {
+  let summary = "loop yield and termination operation";
+  let description = [{
+    The `coredsl.yield` operation is equivalent to the `scf.yield` operation,
+    but only works for `coredsl.switch`
+    If `coredsl.yield` has any operands, the operands must match the parent
+    operation's results.
+    If the parent operation defines no values, then the `coredsl.yield` may be
+    left out in the custom syntax and the builders will insert one implicitly.
+    Otherwise, it has to be present in the syntax to indicate which values are
+    yielded.
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$results);
+  let builders = [OpBuilder<(ins), [{ /* nothing to do */ }]>];
+
+  let assemblyFormat =
+      [{  attr-dict ($results^ `:` type($results))? }];
+}
+
+def SwitchOp : CoreDSL_Op<"switch", [RecursiveMemoryEffects,
+    SingleBlockImplicitTerminator<"coredsl::YieldOp">,
+    DeclareOpInterfaceMethods<RegionBranchOpInterface,
+                              ["getRegionInvocationBounds",
+                               "getEntrySuccessorRegions"]>]> {
+  let summary = "switch-case operation on arbitrary integers";
+  let description = [{
+    `coredsl.switch` is a modified version of `scf.index_switch` that branches
+    to one of the given regions based on the values of the argument and the
+    cases. As opposed to `scf.index_switch`, this operation accepts any integer
+    as argument.
+
+    The operation always has a "default" region and any number of case regions
+    denoted by integer constants. Control-flow transfers to the case region
+    whose constant value equals the value of the argument. If the argument does
+    not equal any of the case values, control-flow transfer to the "default"
+    region.
+
+    Example:
+
+    ```mlir
+    %0 = coredsl.switch %arg0 : ui32 -> ui32
+    case 2 {
+      %1 = hwarith.constant 10 : ui32
+      coredsl.yield %1 : ui32
+    }
+    case 5 {
+      %2 = hwarith.constant 20 : ui32
+      coredsl.yield %2 : ui32
+    }
+    default {
+      %3 = hwarith.constant 30 : ui32
+      coredsl.yield %3 : ui32
+    }
+    ```
+  }];
+
+  let arguments = (ins AnyInteger:$arg, ArrayAttr:$cases);
+  let results = (outs Variadic<AnyType>:$results);
+  let regions = (region SizedRegion<1>:$defaultRegion,
+                        VariadicRegion<SizedRegion<1>>:$caseRegions);
+
+  let assemblyFormat = [{
+    $arg attr-dict `:` type($arg) (`->` type($results)^)?
+    custom<SwitchCases>($cases, $caseRegions) `\n`
+    `` `default` $defaultRegion
+  }];
+
+  let extraClassDeclaration = [{
+    /// Get the number of cases.
+    unsigned getNumCases();
+
+    /// Get the default region body.
+    Block &getDefaultBlock();
+
+    /// Get the body of a case region.
+    Block &getCaseBlock(unsigned idx);
+  }];
+
+  let hasCanonicalizer = 1;
+  let hasVerifier = 1;
+}
+
 #endif // SHORTNAIL_DIALECT_COREDSL_COREDSLOPS_TD

lib/Conversion/CoreDSLToPy/CoreDSLToPy.cpp

@@ -153,8 +153,26 @@ handleWordWiseAccess(Op op,
   }
 }
 
+// Handles both scf::YieldOp and coredsl::YieldOp, as their handling is
+// identical
+template <typename YieldOp>
+static WalkResult emitYieldOp(mlir::raw_indented_ostream &os, YieldOp yieldOp) {
+  os << "return ";
+  bool first = true;
+  for (auto res : yieldOp.getResults()) {
+    if (!first)
+      os << ", ";
+
+    valToPy(os, res);
+    first = false;
+  }
+  os << "\n";
+  return WalkResult::advance();
+}
+
 static WalkResult emitCoreDSLOp(mlir::raw_indented_ostream &os, Operation *op) {
   assert(isa<coredsl::CoreDSLDialect>(op->getDialect()));
+  static unsigned uniqueSwitchNumber = 0;
 
   return TypeSwitch<Operation *, WalkResult>(op)
       .Case<coredsl::SetOp>([&](auto setOp) {
@@ -468,6 +486,86 @@ static WalkResult emitCoreDSLOp(mlir::raw_indented_ostream &os, Operation *op) {
            << (targetType.isSigned() ? "True" : "False") << ")\n";
         return WalkResult::advance();
       })
+      .Case<coredsl::SwitchOp>([&](coredsl::SwitchOp switchOp) {
+        const unsigned switchId = uniqueSwitchNumber++;
+        for (unsigned i = 0, end = switchOp.getNumCases(); i < end; ++i) {
+          Block &caseBlock = switchOp.getCaseBlock(i);
+          // Emit each case as a helper function. This way each
+          // switch has the form
+          // 'ret1, ret2, ..., retn = helper_function', thereby
+          // making handling of the results easier and allowing
+          // us to just emit a return for every yield op
+          os << "def helper_function_coredsl_switch_" << switchId << "_case_"
+             << i << "():\n";
+          os.indent();
+          auto res = caseBlock.template walk<WalkOrder::PreOrder>(
+              [&os](Operation *op) { return emitOp(os, op); });
+          if (res.wasInterrupted())
+            return res;
+          os.unindent();
+        }
+        os << "def helper_function_coresl_switch_" << switchId
+           << "_default_case():\n";
+        os.indent();
+        auto res = switchOp.getDefaultBlock().walk<WalkOrder::PreOrder>(
+            [&os](Operation *op) { return emitOp(os, op); });
+        if (res.wasInterrupted())
+          return res;
+        os.unindent();
+        bool first = true;
+        for (const auto &[idx, attr] : llvm::enumerate(switchOp.getCases())) {
+          const IntegerAttr &intAttr = cast<IntegerAttr>(attr);
+          const APInt &caseValAPInt = intAttr.getValue();
+          const unsigned bitWidth = switchOp.getArg().getType().getWidth();
+          const auto signedness = switchOp.getArg().getType().getSignedness();
+          assert(signedness != IntegerType::Signless);
+          const APSInt caseVal = APSInt(caseValAPInt.sextOrTrunc(bitWidth),
+                                        signedness == IntegerType::Signed);
+          if (first)
+            os << "if ";
+          else
+            os << "elif ";
+          first = false;
+          valToPy(os, switchOp.getArg());
+          os << "." << hwarith::ICmpPredicate::eq << "(";
+          os << intToPy(caseVal) << "):\n";
+          os.indent();
+          if (switchOp.getNumResults() > 0) {
+            bool first = true;
+            for (auto res : switchOp.getResults()) {
+              if (!first) {
+                os << ", ";
+              }
+              first = false;
+              valToPy(os, res);
+            }
+            os << " = ";
+          }
+          os << "helper_function_coredsl_switch_" << switchId << "_case_" << idx
+             << "()\n";
+          os.unindent();
+        }
+        os << "else:\n";
+        os.indent();
+        if (switchOp.getNumResults() > 0) {
+          bool first = true;
+          for (auto res : switchOp.getResults()) {
+            if (!first) {
+              os << ", ";
+            }
+            first = false;
+            valToPy(os, res);
+          }
+          os << " = ";
+        }
+        os << "helper_function_coredsl_switch_" << switchId
+           << "_default_case()\n";
+        os.unindent();
+        // Skip the child scopes, as we already visited them
+        return WalkResult::skip();
+      })
+      .Case<coredsl::YieldOp>(
+          [&](coredsl::YieldOp yieldOp) { return emitYieldOp(os, yieldOp); })
       .Default([&](auto _) {
         op->emitError() << "CoreDSLToPy::emitCoreDSLOp lacks emission code for "
                            "this operation!";
@@ -646,19 +744,8 @@ static WalkResult emitSCFOp(mlir::raw_indented_ostream &os, Operation *op) {
 
         return WalkResult::skip();
       })
-      .Case<scf::YieldOp>([&](auto yieldOp) {
-        os << "return ";
-        bool first = true;
-        for (auto res : yieldOp.getResults()) {
-          if (!first)
-            os << ", ";
-
-          valToPy(os, res);
-          first = false;
-        }
-        os << "\n";
-        return WalkResult::advance();
-      })
+      .Case<scf::YieldOp>(
+          [&](auto yieldOp) { return emitYieldOp(os, yieldOp); })
       // .Case<scf::WhileOp>([&](auto whileOp) {
       //   // TODO PITA
       //   return WalkResult::interrupt();