[Mlir-commits] [mlir] [mlir][ArmSME] Replace nested-region assertion in tile allocation with diagnostic (PR #181934)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-sme

Author: Ayush Kumar Gaur (Ayush3941)

<details>
<summary>Changes</summary>

### What this PR does  

Replace the nested-region assertion in ArmSME tile allocation with a proper diagnostic and graceful failure.

### Why its needed

generateOperationNumbering assumes flattened control flow and currently asserts when encountering ArmSME tile ops inside nested regions. Running the pass standalone on valid IR with structured control flow causes a crash instead of reporting a user-facing error.

### Whats the Fix

Detect nested ArmSME operations during numbering, emit an error indicating that flattened CFG is required (e.g. run -convert-scf-to-cf), and return failure() so the pass terminates cleanly. Add tests verifying diagnostic emission rather than assertion failure.

---
Full diff: https://github.com/llvm/llvm-project/pull/181934.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp (+20-6) 
- (added) mlir/test/Dialect/ArmSME/tile-allocation-nested-regions.mlir (+39) 


``````````diff
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
index 3868e45f34225..73ef150301e36 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
@@ -52,6 +52,7 @@
 #include "mlir/Dialect/ArmSME/Transforms/Transforms.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/IntervalMap.h"
 #include "llvm/ADT/SmallVectorExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
@@ -343,24 +344,34 @@ struct LiveRange {
 /// Operations are numbered consecutively wihin blocks, and the blocks are
 /// topologically sorted (using forward edges). This function is only correct if
 /// all ArmSME have been converted to CF (which is asserted).
-DenseMap<Operation *, unsigned>
+FailureOr<DenseMap<Operation *, unsigned>>
 generateOperationNumbering(FunctionOpInterface function) {
   unsigned index = 0;
   SetVector<Block *> blocks =
       getBlocksSortedByDominance(function.getFunctionBody());
   DenseMap<Operation *, unsigned> operationToIndexMap;
+  bool sawNestedArmSMEOp = false;
   for (Block *block : blocks) {
     index++; // We want block args to have their own number.
     for (Operation &op : block->getOperations()) {
-#ifndef NDEBUG
       op.walk([&](ArmSMETileOpInterface nestedOp) {
-        assert(&op == nestedOp.getOperation() &&
-               "ArmSME tile allocation does not support nested regions");
+        if (&op == nestedOp.getOperation())
+          return;
+        nestedOp.emitError(
+            "ArmSME tile allocation requires flattened control flow; run "
+            "-convert-scf-to-cf before this pass (e.g. via "
+            "convert-arm-sme-to-llvm pipeline)");
+        sawNestedArmSMEOp = true;
       });
-#endif
+      if (sawNestedArmSMEOp)
+        break;
       operationToIndexMap.try_emplace(&op, index++);
     }
+    if (sawNestedArmSMEOp)
+      break;
   }
+  if (sawNestedArmSMEOp)
+    return failure();
   return operationToIndexMap;
 }
 
@@ -809,7 +820,10 @@ LogicalResult mlir::arm_sme::allocateSMETiles(FunctionOpInterface function,
 
   // 2. Gather live ranges for each ArmSME tile within the function.
   Liveness liveness(function);
-  auto operationToIndexMap = generateOperationNumbering(function);
+  auto operationToIndexMapOr = generateOperationNumbering(function);
+  if (failed(operationToIndexMapOr))
+    return failure();
+  auto &operationToIndexMap = *operationToIndexMapOr;
   auto initialLiveRanges = gatherTileLiveRanges(
       operationToIndexMap, liveRangeAllocator, liveness, function);
   if (initialLiveRanges.empty())
diff --git a/mlir/test/Dialect/ArmSME/tile-allocation-nested-regions.mlir b/mlir/test/Dialect/ArmSME/tile-allocation-nested-regions.mlir
new file mode 100644
index 0000000000000..ec8fb98a778b2
--- /dev/null
+++ b/mlir/test/Dialect/ArmSME/tile-allocation-nested-regions.mlir
@@ -0,0 +1,39 @@
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(test-arm-sme-tile-allocation))" -split-input-file -verify-diagnostics
+
+module {
+  func.func @arm_sme_tile_load_hor_f16(%arg0: memref<?x?xf16>) {
+    %c0 = arith.constant 0 : index
+    %0 = arm_sme.get_tile : vector<[8]x[8]xf16>
+    %c8 = arith.constant 8 : index
+    %vscale = vector.vscale
+    %c8_vscale = arith.muli %c8, %vscale : index
+    %cst = arith.constant dense<true> : vector<[8]xi1>
+    %c0_0 = arith.constant 0 : index
+    %c64 = arith.constant 64 : index
+
+    %1 = scf.for %arg1 = %c0_0 to %c64 step %c8 iter_args(%arg2 = %0) -> (vector<[8]x[8]xf16>) {
+      %2 = arith.addi %arg1, %c8 : index
+      // expected-error @+1 {{ArmSME tile allocation requires flattened control flow; run -convert-scf-to-cf before this pass}}
+      %3 = arm_sme.load_tile_slice %arg0[%2, %arg1], %cst, %arg2, %arg1
+           : memref<?x?xf16>, vector<[8]xi1>, vector<[8]x[8]xf16>
+      scf.yield %3 : vector<[8]x[8]xf16>
+    }
+    return
+  }
+}
+
+// -----
+
+module {
+  func.func @main() {
+    %0 = index.constant 0
+    %1 = arm_sme.get_tile : vector<[8]x[8]xi16>
+
+    %2 = scf.index_switch %0 -> vector<[8]x[8]xi16> default {
+      // expected-error @+1 {{ArmSME tile allocation requires flattened control flow; run -convert-scf-to-cf before this pass}}
+      %3 = arm_sme.get_tile : vector<[8]x[8]xi16>
+      scf.yield %3 : vector<[8]x[8]xi16>
+    }
+    return
+  }
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/181934