[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:08:58 PST 2026
https://github.com/Ayush3941 updated https://github.com/llvm/llvm-project/pull/181934
>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/3] [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/3] [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>
}
>From fa3fcd2c2a62d5e0137fe572942a126ad5ca1cf8 Mon Sep 17 00:00:00 2001
From: Ayush3941 <ayushkgaur1 at gmail.com>
Date: Tue, 17 Feb 2026 18:08:23 -0500
Subject: [PATCH 3/3] [mlir][ArmSME] Fix TileAllocation transform regression
-v3
---
mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
index 73ef150301e36..54f142b46510d 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
@@ -52,7 +52,6 @@
#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"
More information about the Mlir-commits
mailing list