[mlir] Allow data flow analysis of non-control flow branch arguments

[krzysz00]

Summary

This commit adds the visitNonControlFlowArguments method to
DataFlowAnalysis, allowing analyses to provide lattice values for the
arguments to a RegionSuccessor block that aren't directly tied to an
op's inputs. For example, integer range interface can use this method
to infer bounds for the step values in loops.

This method has a default implementation that keeps the old behavior
of assigning a pessimistic fixedpoint state to all such arguments.

[Mogball]

mlir/include/mlir/Analysis/DataFlowAnalysis.h

368

This should be moved to the source file.

368

Name this one successor and regionPtr region

[rriddle]

Can you please add a test case that exercises this?

[krzysz00]

@rriddle Do you have thoughts on what such a testcase would look like? The method overall is exercised by the integer range analysis, and the only other available dataflow analysis is SCCP, which doesn't look into region arguments at all (from the looks of it)

[Mogball]

You can write a dummy test pass that just prints some text.

[krzysz00]

Add test for data flow pass and default branch argument methods

[krzysz00]

I've gone and added the test, hopefully this'll do it

mlir/include/mlir/Analysis/DataFlowAnalysis.h

368

This can't be moved to the source file - templates.

[Mogball]

mlir/include/mlir/Analysis/DataFlowAnalysis.h

257–260

The default implementation comes later

368

ah, right you are

370

377

mlir/lib/Analysis/DataFlowAnalysis.cpp

524

mlir/test/lib/Analysis/TestDataFlow.cpp

24

move this to the bottom

39

46

why is this necessary?

[krzysz00]

Address review comments

[krzysz00]

mlir/test/lib/Analysis/TestDataFlow.cpp

46

The base class has a virtual destructor, is why I think this is here.

I was following the lead of the SCCP pass on this.

[Mogball]

mlir/include/mlir/Analysis/DataFlowAnalysis.h

253–257

And can you make the same edits to the copy-pasted doc in the default implementation below?

mlir/lib/Analysis/DataFlowAnalysis.cpp

526–529

Could you add an assert that the lattice values of block arguments that aren't aliases by control flow have values to ensure implementation correctness?

mlir/test/lib/Analysis/TestDataFlow.cpp

46

I'm 90% sure this isn't needed :O

79

Can you implement visitNonControlFlowArguments in the test pass and do something different?

[rriddle]

LG with Jeff's comments resolved.

mlir/lib/Analysis/DataFlowAnalysis.cpp

583–589

Can you split this functor out into a separate variable? the formatting here it quite weird.

[krzysz00]

Address review feedback, fire off tests before landing

mlir/include/mlir/Analysis/DataFlowAnalysis.h

public:
   /// non-null), compute a specific set of region successors that would be
   /// selected.
   virtual void
   getSuccessorsForOperands(RegionBranchOpInterface branch,
                            Optional<unsigned> sourceIndex,
                            ArrayRef<AbstractLatticeElement *> operands,
                            SmallVectorImpl<RegionSuccessor> &successors) = 0;
 
+  /// Given a operation with successor regions, one of those regions,
+  /// and the lattice elements corresponding to the operation's
+  /// arguments, compute the latice values for block arguments
+  /// that are not accounted for by the branching control flow (ex. the
+  /// bounds of loops).
+  virtual ChangeResult
+  visitNonControlFlowArguments(Operation *op, const RegionSuccessor &region,
+                               ArrayRef<AbstractLatticeElement *> operands) = 0;
+
   /// Create a new uninitialized lattice element. An optional value is provided
   /// which, if valid, should be used to initialize the known conservative state
   /// of the lattice.
   virtual AbstractLatticeElement *createLatticeElement(Value value = {}) = 0;
 
 private:
   /// A map from SSA value to lattice element.
   DenseMap<Value, AbstractLatticeElement *> latticeValues;
protected:
   getSuccessorsForOperands(RegionBranchOpInterface branch,
                            Optional<unsigned> sourceIndex,
                            ArrayRef<LatticeElement<ValueT> *> operands,
                            SmallVectorImpl<RegionSuccessor> &successors) {
     SmallVector<Attribute> constantOperands(operands.size());
     branch.getSuccessorRegions(sourceIndex, constantOperands, successors);
   }
 
+  /// Given a operation with successor regions, one of those regions,
+  /// and the lattice elements corresponding to the operation's
+  /// arguments, compute the latice values for block arguments
+  /// that are not accounted for by the branching control flow (ex. the
+  /// bounds of loops). By default, this method marks all such lattice elements
+  /// as having reached a pessimistic fixpoint. The region in the
+  /// RegionSuccessor and the operand latice elements are guaranteed to be
+  /// non-null.
+  virtual ChangeResult
+  visitNonControlFlowArguments(Operation *op, const RegionSuccessor &successor,
+                               ArrayRef<LatticeElement<ValueT> *> operands) {
+    ChangeResult result = ChangeResult::NoChange;
+    Region *region = successor.getSuccessor();
+    ValueRange succArgs = successor.getSuccessorInputs();
+    Block *block = &region->front();
+    Block::BlockArgListType arguments = block->getArguments();
+    if (arguments.size() != succArgs.size()) {
+      unsigned firstArgIdx =
+          succArgs.empty() ? 0
+                           : succArgs[0].cast<BlockArgument>().getArgNumber();
+      result |= markAllPessimisticFixpoint(arguments.take_front(firstArgIdx));
+      result |= markAllPessimisticFixpoint(
+          arguments.drop_front(firstArgIdx + succArgs.size()));
+    }
+    return result;
+  }
+
 private:
   /// Type-erased wrappers that convert the abstract lattice operands to derived
   /// lattices and invoke the virtual hooks operating on the derived lattices.
   ChangeResult
   visitOperation(Operation *op,
                  ArrayRef<detail::AbstractLatticeElement *> operands) final {
     LatticeElement<ValueT> *const *derivedOperandBase =
         reinterpret_cast<LatticeElement<ValueT> *const *>(operands.data());
getSuccessorsForOperands(RegionBranchOpInterface branch,
                            ArrayRef<detail::AbstractLatticeElement *> operands,
                            SmallVectorImpl<RegionSuccessor> &successors) final {
     LatticeElement<ValueT> *const *derivedOperandBase =
         reinterpret_cast<LatticeElement<ValueT> *const *>(operands.data());
     getSuccessorsForOperands(
         branch, sourceIndex,
         llvm::makeArrayRef(derivedOperandBase, operands.size()), successors);
   }
+  ChangeResult visitNonControlFlowArguments(
+      Operation *op, const RegionSuccessor &region,
+      ArrayRef<detail::AbstractLatticeElement *> operands) final {
+    LatticeElement<ValueT> *const *derivedOperandBase =
+        reinterpret_cast<LatticeElement<ValueT> *const *>(operands.data());
+    return visitNonControlFlowArguments(
+        op, region, llvm::makeArrayRef(derivedOperandBase, operands.size()));
+  }
 
   /// Create a new uninitialized lattice element. An optional value is provided,
   /// which if valid, should be used to initialize the known conservative state
   /// of the lattice.
   detail::AbstractLatticeElement *createLatticeElement(Value value) final {
     ValueT knownValue = value ? ValueT::getPessimisticValueState(value)
                               : ValueT::getPessimisticValueState(context);
     return new (allocator.Allocate()) LatticeElement<ValueT>(knownValue);

mlir/lib/Analysis/DataFlowAnalysis.cpp

 //===- DataFlowAnalysis.cpp -----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/DataFlowAnalysis.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallPtrSet.h"
 
 #include <queue>
 
 using namespace mlir;
 using namespace mlir::detail;
 
 namespace {
void visitRegionBranchOperation(
       ArrayRef<AbstractLatticeElement *> operandLattices);
 
   /// Visit the given set of region successors, computing any necessary lattice
   /// state. The provided function returns the input operands to the region at
   /// the given index. If the index is 'None', the input operands correspond to
   /// the parent operation results.
   void visitRegionSuccessors(
       Operation *parentOp, ArrayRef<RegionSuccessor> regionSuccessors,
+      ArrayRef<AbstractLatticeElement *> operandLattices,
       function_ref<OperandRange(Optional<unsigned>)> getInputsForRegion);
 
   /// Visit the given terminator operation and compute any necessary lattice
   /// state.
   void
   visitTerminatorOperation(Operation *op,
                            ArrayRef<AbstractLatticeElement *> operandLattices);
 
void ForwardDataFlowSolver::visitRegionBranchOperation(
   // any results of this operation as overdefined, as we can't reason about
   // them.
   // TODO: If we had an interface to detect pass through operands, we could
   // resolve some results based on the lattice state of the operands. We could
   // also allow for the parent operation to have itself as a region successor.
   if (successors.empty())
     return markAllPessimisticFixpoint(branch, branch->getResults());
   return visitRegionSuccessors(
-      branch, successors, [&](Optional<unsigned> index) {
+      branch, successors, operandLattices, [&](Optional<unsigned> index) {
         assert(index && "expected valid region index");
         return branch.getSuccessorEntryOperands(*index);
       });
 }
 
 void ForwardDataFlowSolver::visitRegionSuccessors(
     Operation *parentOp, ArrayRef<RegionSuccessor> regionSuccessors,
+    ArrayRef<AbstractLatticeElement *> operandLattices,
     function_ref<OperandRange(Optional<unsigned>)> getInputsForRegion) {
   for (const RegionSuccessor &it : regionSuccessors) {
     Region *region = it.getSuccessor();
     ValueRange succArgs = it.getSuccessorInputs();
 
     // Check to see if this is the parent operation.
     if (!region) {
       ResultRange results = parentOp->getResults();
for (const RegionSuccessor &it : regionSuccessors) {
     markBlockExecutable(entryBlock);
 
     // If all of the arguments have already reached a fixpoint, the arguments
     // have already been fully resolved.
     Block::BlockArgListType arguments = entryBlock->getArguments();
     if (llvm::all_of(arguments, [&](Value arg) { return isAtFixpoint(arg); }))
       continue;
 
-    // Mark any arguments that do not receive inputs as having reached a
-    // pessimistic fixpoint, we won't be able to discern if they are constant.
-    // TODO: This isn't exactly ideal. There may be situations in which a
-    // region operation can provide information for certain results that
-    // aren't part of the control flow.
     if (succArgs.size() != arguments.size()) {
-      if (succArgs.empty()) {
-        markAllPessimisticFixpoint(arguments);
-        continue;
+      if (analysis.visitNonControlFlowArguments(
+              parentOp, it, operandLattices) == ChangeResult::Change) {
+        unsigned firstArgIdx =
+            succArgs.empty() ? 0
+                             : succArgs[0].cast<BlockArgument>().getArgNumber();
+        for (Value v : arguments.take_front(firstArgIdx)) {
+          assert(!analysis.getLatticeElement(v).isUninitialized() &&
+                 "Non-control flow block arg has no lattice value after "
+                 "analysis callback");
+          visitUsers(v);
+        }
+        for (Value v : arguments.drop_front(firstArgIdx + succArgs.size())) {
+          assert(!analysis.getLatticeElement(v).isUninitialized() &&
+                 "Non-control flow block arg has no lattice value after "
+                 "analysis callback");
+          visitUsers(v);
+        }
       }
-
-      unsigned firstArgIdx = succArgs[0].cast<BlockArgument>().getArgNumber();
-      markAllPessimisticFixpointAndVisitUsers(
-          arguments.take_front(firstArgIdx));
-      markAllPessimisticFixpointAndVisitUsers(
-          arguments.drop_front(firstArgIdx + succArgs.size()));
     }
 
     // Update the lattice of arguments that have inputs from the predecessor.
     OperandRange succOperands = getInputsForRegion(region->getRegionNumber());
     for (auto it : llvm::zip(succArgs, succOperands)) {
       AbstractLatticeElement &argValue =
           analysis.getLatticeElement(std::get<0>(it));
       AbstractLatticeElement &operandValue =
analysis.getSuccessorsForOperands(regionInterface,
                                       parentRegion->getRegionNumber(),
                                       operandLattices, regionSuccessors);
     if (regionSuccessors.empty())
       return;
 
     // Try to get "region-like" successor operands if possible in order to
     // propagate the operand states to the successors.
     if (isRegionReturnLike(op)) {
-      return visitRegionSuccessors(
-          parentOp, regionSuccessors, [&](Optional<unsigned> regionIndex) {
-            // Determine the individual region successor operands for the given
-            // region index (if any).
-            return *getRegionBranchSuccessorOperands(op, regionIndex);
-          });
+      auto getOperands = [&](Optional<unsigned> regionIndex) {
+        // Determine the individual region  successor operands for the given
+      };
+      return visitRegionSuccessors(parentOp, regionSuccessors, operandLattices,
+                                   getOperands);
     }
 
     // If this terminator is not "region-like", conservatively mark all of the
     // successor values as having reached the pessimistic fixpoint.
     for (auto &it : regionSuccessors) {
       // If the successor is a region, mark the entry block as executable so
       // that we visit operations defined within. If the successor is the
       // parent operation, we simply mark the control flow results as having
       // reached the pessimistic state.

mlir/test/Analysis/test-data-flow.mlir

+// RUN: mlir-opt -test-data-flow --allow-unregistered-dialect %s 2>&1 | FileCheck %s
+
+// CHECK-LABEL: Testing : "loop-arg-pessimistic"
+module attributes {test.name = "loop-arg-pessimistic"} {
+  func @f() -> index {
+    // CHECK: Visiting : %{{.*}} = arith.constant 0
+    // CHECK-NEXT: Result 0 moved from uninitialized to 1
+    %c0 = arith.constant 0 : index
+    // CHECK: Visiting : %{{.*}} = arith.constant 1
+    // CHECK-NEXT: Result 0 moved from uninitialized to 1
+    %c1 = arith.constant 1 : index
+    // CHECK: Visiting region branch op : %{{.*}} = scf.for
+    // CHECK: Block argument 0 moved from uninitialized to 1
+    %0 = scf.for %arg1 = %c0 to %c1 step %c1 iter_args(%arg2 = %c0) -> index {
+      // CHECK: Visiting : %{{.*}} = arith.addi %{{.*}}, %{{.*}}
+      // CHECK-NEXT: Arg 0 : 1
+      // CHECK-NEXT: Arg 1 : 1
+      // CHECK-NEXT: Result 0 moved from uninitialized to 1
+      %10 = arith.addi %arg1, %arg2 : index
+      scf.yield %10 : index
+    }
+    return %0 : index
+  }
+}

mlir/test/lib/Analysis/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestAnalysis
   TestAliasAnalysis.cpp
   TestCallGraph.cpp
+  TestDataFlow.cpp
   TestLiveness.cpp
   TestMatchReduction.cpp
   TestMemRefBoundCheck.cpp
   TestMemRefDependenceCheck.cpp
   TestMemRefStrideCalculation.cpp
   TestSlice.cpp
 
 

mlir/test/lib/Analysis/TestDataFlow.cpp

+//===- TestDataFlow.cpp - Test data flow analysis system -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains test passes for defining and running a dataflow analysis.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/DataFlowAnalysis.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/STLExtras.h"
+
+using namespace mlir;
+
+namespace {
+struct WasAnalyzed {
+  WasAnalyzed(bool wasAnalyzed) : wasAnalyzed(wasAnalyzed) {}
+
+  static WasAnalyzed join(const WasAnalyzed &a, const WasAnalyzed &b) {
+    return a.wasAnalyzed && b.wasAnalyzed;
+  }
+
+  static WasAnalyzed getPessimisticValueState(MLIRContext *context) {
+    return false;
+  }
+
+  static WasAnalyzed getPessimisticValueState(Value v) {
+    return getPessimisticValueState(v.getContext());
+  }
+
+  bool operator==(const WasAnalyzed &other) const {
+    return wasAnalyzed == other.wasAnalyzed;
+  }
+
+  bool wasAnalyzed;
+};
+
+struct TestAnalysis : public ForwardDataFlowAnalysis<WasAnalyzed> {
+  using ForwardDataFlowAnalysis<WasAnalyzed>::ForwardDataFlowAnalysis;
+
+  ChangeResult
+  visitOperation(Operation *op,
+                 ArrayRef<LatticeElement<WasAnalyzed> *> operands) final {
+    ChangeResult ret = ChangeResult::NoChange;
+    llvm::errs() << "Visiting : ";
+    op->print(llvm::errs());
+    llvm::errs() << "\n";
+
+    WasAnalyzed result(true);
+    for (auto &pair : llvm::enumerate(operands)) {
+      LatticeElement<WasAnalyzed> *elem = pair.value();
+      llvm::errs() << "Arg " << pair.index();
+      if (!elem->isUninitialized()) {
+        llvm::errs() << " : " << elem->getValue().wasAnalyzed << "\n";
+        result = WasAnalyzed::join(result, elem->getValue());
+      } else {
+        llvm::errs() << " uninitialized\n";
+      }
+    }
+    for (const auto &pair : llvm::enumerate(op->getResults())) {
+      LatticeElement<WasAnalyzed> &lattice = getLatticeElement(pair.value());
+      llvm::errs() << "Result " << pair.index() << " moved from ";
+      if (lattice.isUninitialized())
+        llvm::errs() << "uninitialized";
+      else
+        llvm::errs() << lattice.getValue().wasAnalyzed;
+      ret |= lattice.join({result});
+      llvm::errs() << " to " << lattice.getValue().wasAnalyzed << "\n";
+    }
+    return ret;
+  }
+
+  ChangeResult visitNonControlFlowArguments(
+      Operation *op, const RegionSuccessor &successor,
+      ArrayRef<LatticeElement<WasAnalyzed> *> operands) final {
+    ChangeResult ret = ChangeResult::NoChange;
+    llvm::errs() << "Visiting region branch op : ";
+    op->print(llvm::errs());
+    llvm::errs() << "\n";
+
+    Region *region = successor.getSuccessor();
+    Block *block = &region->front();
+    Block::BlockArgListType arguments = block->getArguments();
+    // Mark all arguments to blocks as analyzed unless they already have
+    // an unanalyzed state.
+    for (const auto &pair : llvm::enumerate(arguments)) {
+      LatticeElement<WasAnalyzed> &lattice = getLatticeElement(pair.value());
+      llvm::errs() << "Block argument " << pair.index() << " moved from ";
+      if (lattice.isUninitialized())
+        llvm::errs() << "uninitialized";
+      else
+        llvm::errs() << lattice.getValue().wasAnalyzed;
+      ret |= lattice.join({true});
+      llvm::errs() << " to " << lattice.getValue().wasAnalyzed << "\n";
+    }
+    return ret;
+  }
+};
+
+struct TestDataFlowPass
+    : public PassWrapper<TestDataFlowPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDataFlowPass)
+
+  StringRef getArgument() const final { return "test-data-flow"; }
+  StringRef getDescription() const final {
+    return "Print the actions taken during a dataflow analysis.";
+  }
+  void runOnOperation() override {
+    llvm::errs() << "Testing : " << getOperation()->getAttr("test.name")
+                 << "\n";
+    TestAnalysis analysis(getOperation().getContext());
+    analysis.run(getOperation());
+  }
+};
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestDataFlowPass() { PassRegistration<TestDataFlowPass>(); }
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestAffineLoopParametricTilingPass();
 void registerTestAliasAnalysisPass();
 void registerTestBuiltinAttributeInterfaces();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestControlFlowSink();
 void registerTestGpuSerializeToCubinPass();
 void registerTestGpuSerializeToHsacoPass();
+void registerTestDataFlowPass();
 void registerTestDataLayoutQuery();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDiagnosticsPass();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandTanhPass();
 void registerTestComposeSubView();
 void registerTestMultiBuffering();
void registerTestPasses() {
   mlir::test::registerTestDiagnosticsPass();
 #if MLIR_CUDA_CONVERSIONS_ENABLED
   mlir::test::registerTestGpuSerializeToCubinPass();
 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED
   mlir::test::registerTestGpuSerializeToHsacoPass();
 #endif
   mlir::test::registerTestDecomposeCallGraphTypes();
+  mlir::test::registerTestDataFlowPass();
   mlir::test::registerTestDataLayoutQuery();
   mlir::test::registerTestDominancePass();
   mlir::test::registerTestDynamicPipelinePass();
   mlir::test::registerTestExpandTanhPass();
   mlir::test::registerTestComposeSubView();
   mlir::test::registerTestMultiBuffering();
   mlir::test::registerTestGpuParallelLoopMappingPass();
   mlir::test::registerTestIRVisitorsPass();