Request for advice in updating IntRange optimziations

MLIR
1

In our out-of-tree dialects, we have region-holding ops with special semantics, in particular, a region can yield an integer or tensor-of-ints to its parent op via control flow, and in the process the integer type is wrapped in a different type:

      func.func @repro_crash() -> !secret.secret<tensor<32xi16>> {
        %cst = arith.constant dense<1> : tensor<32xi16>
        %0 = secret.generic() {
          %1 = tensor_ext.assign_layout %cst {layout = ...} : tensor<32xi16>
          secret.yield %1 : tensor<32xi16>
        } -> (!secret.secret<tensor<32xi16>> {tensor_ext.layout = #layout})
        return %0 : !secret.secret<tensor<32xi16>>
      }

If I run IntegerRangeAnalysis on this and then applyPatternsGreedily using arith::populateIntRangeOptimizationsPatterns for the pattern set, the resulting run stack traces.

In particular, the int range propagates to the result of the secret.generic op (which implements RegionBranchOpInterface), but then when MaterializeKnownConstantValues executes, we hit these two forced options for the materialization: llvm-project/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp at 2296100e3b7ec901d972e88785768d3a2c494d72 · llvm/llvm-project · GitHub

constAttr = rewriter.getIntegerAttr(type, *maybeConstValue) gives a stack trace if type is not integer or index.

I’m wondering the right way to fix this upstream. I suspect I should add a check to needsReplacing to ensure that the value has an appropriate type? Or should this be fixed in the integer range analysis more directly (though the analysis itself relies on SparseForwardDataFlowAnalysis which doesn’t give control over how values propagate through RegionBranchOpInterface…)?

2

I can’t offer too much advice here, but you might want to verify if secret.yield implements the RegionBranchTerminatorOpInterface. The dataflow analysis framework relies on it as well.

3

Indeed, it has ReturnLike which as I understand provides a default implementation of the methods for RegionBranchTerminatorOpInterface

4

So I think, with the code as written, you may not be ReturnLike - or at least you may wont to have a more complex definition of RegionBranchOpInterface

Dataflow analysis has default handling for things that look like iteration arguments - values that come in to the region and get yielded/returned/… from it. Integer range analysis handles those by way of the dataflow solver’s general handling.

What you have with secret.yield is what the dataflow analysis framework calls a “non-control flow argument” - something analogous to a loop IV. It’s something that doesn’t flow into the region in some obvious sense.

So you may want to update your interface implementations to make sure that it’s clear that the yielded value isn’t something that should be handled by the usual default handling … since, for one thing, it changes type.

5

It doesn’t flow into the region in that example, but it definitely has a successor: the op result. (In other examples, the generic op does indeed have arguments that flow into the region).

I think the type change is really the issue here: as I understand it, nothing about the dataflow framework requires a value’s type to be equal to the type of its successor, and the IntRangeAnalysis doesn’t override this with a more specific requirement. I think the same issue would occur with say, !async.value when it wraps an integer type.

6

I always expected that these region-like operations needed their control-flow arguments to match types

But yeah, as to the concrete code at issue, I’d sign off on “if the value doesn’t have integer [element] type, fail” as a fix to make your usecase work

7

I think the type change is really the issue here: as I understand it, nothing about the dataflow framework requires a value’s type to be equal to the type of its successor.

My apologies—for some reason, I missed the recent messages in our conversation. Prior to this, I had studied the implementation of the data-flow analysis framework, including the ‘non-control flow arguments’ @krzysz00 mentioned. I agree that this should be achievable, as the framework handles them by default through RegionBranchTerminatorOpInterface and RegionBranchOpInterface.

I also always expected that these region-like operations needed their control-flow arguments to match types.