[Mlir-commits] [mlir] [mlir][ArmSME] Replace nested-region assertion in tile allocation with diagnostic (PR #181934)
Ayush Kumar Gaur
llvmlistbot at llvm.org
Tue Feb 17 15:06:06 PST 2026
https://github.com/Ayush3941 created https://github.com/llvm/llvm-project/pull/181934
### 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.
>From 77b9d84935bbb92e40f1f64203e673b7a6532831 Mon Sep 17 00:00:00 2001
From: Ayush3941 <ayushkgaur1 at gmail.com>
Date: Tue, 17 Feb 2026 17:48:35 -0500
Subject: [PATCH 1/2] [mlir][ArmSME] Fix TileAllocation transform regression
---
.../ArmSME/Transforms/TileAllocation.cpp | 26 ++++++++++----
.../tile-allocation-nested-regions.mlir | 36 +++++++++++++++++++
2 files changed, 56 insertions(+), 6 deletions(-)
create mode 100644 mlir/test/Dialect/ArmSME/tile-allocation-nested-regions.mlir
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..7f44372e0ef03
--- /dev/null
+++ b/mlir/test/Dialect/ArmSME/tile-allocation-nested-regions.mlir
@@ -0,0 +1,36 @@
+// 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
+ // expected-error @+1 {{ArmSME tile allocation requires flattened control flow; run -convert-scf-to-cf before this pass (e.g. via convert-arm-sme-to-llvm pipeline)}}
+ %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
+ %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>
+ // expected-error @+1 {{ArmSME tile allocation requires flattened control flow; run -convert-scf-to-cf before this pass (e.g. via convert-arm-sme-to-llvm pipeline)}}
+ %2 = scf.index_switch %0 -> vector<[8]x[8]xi16> default {
+ %3 = arm_sme.get_tile : vector<[8]x[8]xi16>
+ scf.yield %3 : vector<[8]x[8]xi16>
+ }
+ return
+ }
+}
>From 512b071eaec628c0ce6f2e700fe81bc9175b276d Mon Sep 17 00:00:00 2001
From: Ayush3941 <ayushkgaur1 at gmail.com>
Date: Tue, 17 Feb 2026 18:00:12 -0500
Subject: [PATCH 2/2] [mlir][ArmSME] Fix TileAllocation transform regression
-v2
---
.../Dialect/ArmSME/tile-allocation-nested-regions.mlir | 9 ++++++---
1 file changed, 6 insertions(+), 3 deletions(-)
diff --git a/mlir/test/Dialect/ArmSME/tile-allocation-nested-regions.mlir b/mlir/test/Dialect/ArmSME/tile-allocation-nested-regions.mlir
index 7f44372e0ef03..ec8fb98a778b2 100644
--- a/mlir/test/Dialect/ArmSME/tile-allocation-nested-regions.mlir
+++ b/mlir/test/Dialect/ArmSME/tile-allocation-nested-regions.mlir
@@ -10,10 +10,12 @@ module {
%cst = arith.constant dense<true> : vector<[8]xi1>
%c0_0 = arith.constant 0 : index
%c64 = arith.constant 64 : index
- // expected-error @+1 {{ArmSME tile allocation requires flattened control flow; run -convert-scf-to-cf before this pass (e.g. via convert-arm-sme-to-llvm pipeline)}}
+
%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
- %3 = arm_sme.load_tile_slice %arg0[%2, %arg1], %cst, %arg2, %arg1 : memref<?x?xf16>, vector<[8]xi1>, vector<[8]x[8]xf16>
+ // 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
@@ -26,8 +28,9 @@ module {
func.func @main() {
%0 = index.constant 0
%1 = arm_sme.get_tile : vector<[8]x[8]xi16>
- // expected-error @+1 {{ArmSME tile allocation requires flattened control flow; run -convert-scf-to-cf before this pass (e.g. via convert-arm-sme-to-llvm pipeline)}}
+
%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>
}
More information about the Mlir-commits
mailing list