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

[Sarthak Gupta]

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

>From 45327ec37c75fa6f90464dde88117ffe93f67ebf Mon Sep 17 00:00:00 2001
From: gptsarthak <sarthakgpt95 at gmail.com>
Date: Tue, 10 Oct 2023 18:21:26 +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  | 19 +++++++++++++++++++
 2 files changed, 20 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..d6396af37cb2c38 100644
--- a/mlir/test/Dialect/Affine/parallelize.mlir
+++ b/mlir/test/Dialect/Affine/parallelize.mlir
@@ -323,3 +323,22 @@ func.func @iter_arg_memrefs(%in: memref<10xf32>) {
   }
   return
 }
+
+// CHECK-LABEL: @scf_affine_scope
+func.func @scf_affine_scope() {
+  %c0 = arith.constant 0 : index
+  %0 = tensor.empty(%c0) : tensor<?xi1>
+  %1 = bufferization.to_memref %0 : memref<?xi1>
+  %alloc = memref.alloc(%c0) : memref<?xi1>
+  %2 = scf.index_switch %c0 -> tensor<?x31x6xf16>
+  default {
+    %dim = memref.dim %1, %c0 : memref<?xi1>
+    affine.for %arg0 = 0 to %dim {
+            %3 = affine.load %1[%arg0] : memref<?xi1>
+            affine.store %3, %alloc[%arg0] : memref<?xi1>
+    }
+    %3 = tensor.empty(%c0) : tensor<?x31x6xf16>
+    scf.yield %3 : tensor<?x31x6xf16>
+  } 
+  return
+}