Adding all indices backwardslice to skip

Signed-off-by: jerryyin <[email protected]>
iree-org · Jan 22, 2025 · 23dc093 · 23dc093
1 parent 2e2b1f0
commit 23dc093
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
@@ -2544,14 +2544,10 @@ LogicalResult initGPULaunchConfig(FunctionOpInterface funcOp) {
               dyn_cast<linalg::GenericOp>(indices.getDefiningOp())) {
         genericToSkip.insert(genericOp);
       }
-      // If scatter's backward slices are generic ops, mark them as to skip too.
+      // Mark scatter's backward slices as to skip too.
       SetVector<Operation *> slices;
       getBackwardSlice(indices, &slices);
-      for (auto slice : slices) {
-        if (isa<linalg::GenericOp>(slice)) {
-          genericToSkip.insert(slice);
-        }
-      }
+      genericToSkip.insert(slices.begin(), slices.end());
     }
   }