[Mlir-commits] [mlir] [mlir][affine] Add `AffineScope` trait to `scf.index_switch` (PR #68672)

[Sarthak Gupta]

https://github.com/gptsarthak created https://github.com/llvm/llvm-project/pull/68672

Fixes https://github.com/llvm/llvm-project/issues/64287

The code in AffineStructures.cpp:FlatAffineValueConstraints' assert statement checks that the loop bound (i.e. val) isTopLevelValue or isForInductionVar. In this case, `%dim = memref.dim %20, %c0 : memref<?x?x6xi1>` is neither, since it is inside `scf.index_switch`, which is not a part of `AffineScope`. Thus, the assert is triggered.

However, `%dim = memref.dim %20, %c0 : memref<?x?x6xi1>` is an invariant in the scope of the `scf.index_switch`. If it was a loop,  a invariant code motion could have hoisted it to be outside the `scf` scope, making it a valid `TopLevelValue`.

Adding the `AffineScope` trait to `scf.index_switch` solves the issue. This defines a new top level scope for symbols, which proves highly practical as it enables a broader range of things to be represented as sequences of affine loop nests.
I even suggest that we add this trait to other SCF operations like `scf.for`, `scf.while` as well, as that will be very useful, and should not have any major issues that we cannot solve.


>From 3611f4eb44c16227739a1b3440484f551d2dbefa Mon Sep 17 00:00:00 2001
From: gptsarthak <sarthakgpt95 at gmail.com>
Date: Tue, 10 Oct 2023 13:29:06 +0530
Subject: [PATCH] Add `AffineScope` trait to `scf.index_switch`

Adding the `AffineScope` trait to `scf.index_switch` solves https://github.com/llvm/llvm-project/issues/64287. This defines a new top level scope for symbols, which proves highly practical as it enables a broader range of things to be represented as sequences of affine loop nests.
I suggest that we add this trait to other SCF operations like `scf.for`, `scf.while` as well, as that will be very useful, and should not have any major issues that we cannot solve.
---
 mlir/include/mlir/Dialect/SCF/IR/SCFOps.td |  2 +-
 mlir/test/Dialect/Affine/parallelize.mlir  | 21 +++++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index e1a604a88715f0e..a5e927f660beaa4 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -1101,7 +1101,7 @@ def WhileOp : SCF_Op<"while",
 // IndexSwitchOp
 //===----------------------------------------------------------------------===//
 
-def IndexSwitchOp : SCF_Op<"index_switch", [RecursiveMemoryEffects,
+def IndexSwitchOp : SCF_Op<"index_switch", [RecursiveMemoryEffects, AffineScope,
     SingleBlockImplicitTerminator<"scf::YieldOp">,
     DeclareOpInterfaceMethods<RegionBranchOpInterface,
                               ["getRegionInvocationBounds",
diff --git a/mlir/test/Dialect/Affine/parallelize.mlir b/mlir/test/Dialect/Affine/parallelize.mlir
index b0e121543f2d6f4..9115247abdfb3b0 100644
--- a/mlir/test/Dialect/Affine/parallelize.mlir
+++ b/mlir/test/Dialect/Affine/parallelize.mlir
@@ -323,3 +323,24 @@ func.func @iter_arg_memrefs(%in: memref<10xf32>) {
   }
   return
 }
+
+// CHECK-LABEL: @scf_affine_scope
+func.func @scf_affine_scope(%arg0: memref<9x9xi32>, %arg1: i1) {
+  %idx0 = index.constant 0
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %8 = tensor.empty() : tensor<6xi1>
+  %20 = bufferization.to_memref %8 : memref<6xi1>
+  %alloc_62 = memref.alloc() : memref<6xi1>
+  %26 = scf.index_switch %c1 -> tensor<?x31x6xf16>
+  default {
+    %dim = memref.dim %20, %c0 : memref<6xi1>
+    affine.for %arg2 = 0 to %dim {
+            %38 = affine.load %20[%arg2] : memref<6xi1>
+            affine.store %38, %alloc_62[%arg2] : memref<6xi1>
+    }
+    %38 = tensor.empty(%idx0) : tensor<?x31x6xf16>
+    scf.yield %38 : tensor<?x31x6xf16>
+  } 
+  return
+}