[MLIR][Interfaces] Allow non-builtin scalar types in IndexingMapOpInterface (#188774)

joker-eph · web-flow · commit c5182c192a7d · 2026-03-26T18:22:20.000+01:00
The scalar type check in `IndexingMapOpInterface::verifyImpl` and its helper `verifyIndexingMapOperandType` used `isIntOrIndexOrFloat() || isa<ComplexType>()`, which only accepted builtin scalar types. This rejected valid custom-dialect scalar types such as pointer types (`\!ptr.ptr<...>`) or other non-shaped dialect types. The `isScalar` method in `DestinationStyleOpInterface` already treats any non-MemRef/non-Tensor type as scalar. Align `IndexingMapOpInterface` with this definition by treating any non-ShapedType as a rank-0 scalar, regardless of whether it is a builtin type. Fixes #183606 Assisted-by: Claude Code
diff --git a/mlir/lib/Interfaces/IndexingMapOpInterface.cpp b/mlir/lib/Interfaces/IndexingMapOpInterface.cpp
@@ -16,8 +16,9 @@ namespace mlir {
 
 static LogicalResult verifyIndexingMapOperandType(Operation *op, Type t,
                                                   unsigned operandNumber) {
-  // Scalars are allowed (treated as rank-0). verifyImpl checks the rank.
-  if (t.isIntOrIndexOrFloat() || isa<ComplexType>(t))
+  // Non-shaped types are treated as scalars (rank-0). This includes builtin
+  // types (integer, float, complex) as well as custom dialect types.
+  if (!isa<ShapedType>(t))
     return success();
 
   // Vectors are allowed.
@@ -67,8 +68,9 @@ LogicalResult mlir::IndexingMapOpInterface::verifyImpl() {
     if (indexingMap.getNumSymbols() != 0)
       return this->emitOpError("unexpected symbols in indexing_map #")
              << opOperand.getOperandNumber();
-    // Handle scalars.
-    if (ty.isIntOrIndexOrFloat() || isa<ComplexType>(ty)) {
+    // Handle scalars (non-shaped types: integer, float, complex, custom types,
+    // etc.).
+    if (!isa<ShapedType>(ty)) {
       int64_t rank = 0;
       if (indexingMap.getNumResults() != rank)
         return this->emitOpError("expected operand #")
diff --git a/mlir/test/Dialect/Linalg/fill-custom-scalar-type.mlir b/mlir/test/Dialect/Linalg/fill-custom-scalar-type.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt %s | FileCheck %s
+// This test verifies that linalg.fill and linalg.generic accept non-builtin
+// scalar types (e.g., custom dialect types) as operands.
+
+// CHECK-LABEL: @fill_non_builtin_scalar_type
+func.func @fill_non_builtin_scalar_type(%src: !test.test_type, %dst: tensor<4x!test.test_type>) -> tensor<4x!test.test_type> {
+  // CHECK: linalg.fill
+  %result = linalg.fill ins(%src : !test.test_type) outs(%dst : tensor<4x!test.test_type>) -> tensor<4x!test.test_type>
+  return %result : tensor<4x!test.test_type>
+}
