Integrate llvm-project@681b2413 (#20274)

Retains reverts for 418e07b7e and 1a3dc8e9c.

---------

Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>
Co-authored-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>

Author: qedawkins

compiler/src/iree/compiler/Codegen/Common/BlockDynamicDimensions.cpp

@@ -9,6 +9,7 @@
 #include "iree/compiler/Codegen/Transforms/Transforms.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
 #include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
 #include "iree/compiler/Dialect/Util/IR/UtilOps.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
@@ -233,9 +234,18 @@ static LogicalResult blockDynamicDimensions(
         blockDynamicDimensionsOfValue(rewriter, resultDivisibilityInfo, result);
     if (reshapes) {
       llvm::SmallPtrSet<Operation *, 1> ignoreUses;
-      ignoreUses.insert(reshapes->expandShapeOp);
-      rewriter.replaceAllUsesExcept(
-          result, reshapes->collapseShapeOp.getResult(), ignoreUses);
+      auto replaceIf = [&](OpOperand &use) {
+        Operation *user = use.getOwner();
+        if (user == reshapes->expandShapeOp) {
+          return false;
+        }
+        if (isa<tensor::DimOp>(user)) {
+          return false;
+        }
+        return true;
+      };
+      rewriter.replaceUsesWithIf(result, reshapes->collapseShapeOp.getResult(),
+                                 replaceIf);
     }
   }
   return success();

compiler/src/iree/compiler/Codegen/Common/Passes.td

@@ -23,6 +23,7 @@ def BlockDynamicDimensionsPass
     : Pass<"iree-codegen-block-dynamic-dimensions"> {
   let summary = "Expand dynamic dimensions that are known to be multiples of "
                 "statically known values.";
+  let dependentDialects = ["iree_compiler::IREE::Util::UtilDialect"];
 }
 
 def BubbleUpOrdinalOpsPass : Pass<"iree-codegen-bubble-up-ordinal-ops", ""> {

compiler/src/iree/compiler/Codegen/Common/test/block_dynamic_dims.mlir

@@ -277,3 +277,34 @@ func.func @block_elementwise(%arg0 : tensor<?xf16>, %dim : index)
 //       CHECK:   %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ELEM]]
 //  CHECK-SAME:     tensor<?x1024xf16> into tensor<?xf16>
 //       CHECK:   return %[[COLLAPSE]] : tensor<?xf16>
+
+// -----
+
+// Check that there are no SSA violations during blocking
+func.func @check_ssa_violation(%dim : index,
+    %lhs : tensor<?xf32>, %rhs : tensor<?xf32>) -> tensor<?xf32> {
+  %c4 = arith.constant 4 : index
+  %0 = arith.muli %dim, %c4 overflow<nsw> : index
+  %1 = tensor.empty(%0) : tensor<?xf32>
+  %2 = linalg.generic {
+        indexing_maps = [affine_map<(d0) -> (d0)>,
+                         affine_map<(d0) -> (d0)>,
+                         affine_map<(d0) -> (d0)>],
+        iterator_types = ["parallel"]}
+        ins(%lhs, %rhs : tensor<?xf32>, tensor<?xf32>)
+        outs(%1 : tensor<?xf32>) {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %33 = arith.addf %in, %in_0 : f32
+    linalg.yield %33 : f32
+  } -> tensor<?xf32>
+  return %2 : tensor<?xf32>
+}
+// CHECK-LABEL: func @check_ssa_violation
+//  CHECK-SAME:     %[[LHS:[a-zA-Z0-9]+]]: tensor<?xf32>
+//  CHECK-SAME:     %[[RHS:[a-zA-Z0-9]+]]: tensor<?xf32>
+//   CHECK-DAG:   %[[EXPANDED0:.+]] = tensor.expand_shape %[[LHS]]
+//   CHECK-DAG:   %[[EXPANDED1:.+]] = tensor.expand_shape %[[RHS]]
+//       CHECK:   %[[GENERIC:.+]] = linalg.generic
+//  CHECK-SAME:       ins(%[[EXPANDED0]], %[[EXPANDED1]] :
+//       CHECK:   %[[COLLAPSE:.+]] = tensor.collapse_shape %[[GENERIC]]
+//       CHECK:   return %[[COLLAPSE]]

compiler/src/iree/compiler/Dialect/VM/Conversion/VMToEmitC/ConvertVMToEmitC.cpp

@@ -3170,7 +3170,7 @@ class BranchOpConversion : public EmitCConversionPattern<IREE::VM::BranchOp> {
         }
 
         Value blockArgRef =
-            signatureConversion.getInputMapping(index)->replacementValue;
+            signatureConversion.getInputMapping(index)->replacementValues[0];
 
         assert(isa<IREE::VM::RefType>(operand.getType()));
         assert(isa<emitc::PointerType>(blockArgRef.getType()));

compiler/src/iree/compiler/DispatchCreation/test/bubble_up_expand_shapes.mlir

@@ -77,7 +77,6 @@ util.func public @attention_v_reshape_propagation(%arg0: index,
 util.func @multiple_users(%arg0 : tensor<?x128xf16>,
     %arg1 : tensor<4x?x32x128xf16>) -> tensor<4x?x32x8x16xf16> {
   %c0 = arith.constant 0 : index
-  %c1 = arith.constant 1 : index
   %dim = tensor.dim %arg0, %c0 : tensor<?x128xf16>
   %empty = tensor.empty(%dim) : tensor<4x?x32x128xf16>
   %0 = linalg.generic {
@@ -97,7 +96,7 @@ util.func @multiple_users(%arg0 : tensor<?x128xf16>,
       %2 = arith.addf %1, %b0 : f16
       linalg.yield %2 : f16
   } -> tensor<4x?x32x128xf16>
-  %1 = tensor.dim %0, %c1 : tensor<4x?x32x128xf16>
+  %1 = tensor.dim %arg0, %c0 : tensor<?x128xf16>
   %2 = tensor.expand_shape %0 [[0], [1], [2], [3, 4]] output_shape [4, %1, 32, 8, 16]
       : tensor<4x?x32x128xf16> into tensor<4x?x32x8x16xf16>
   util.return %2 : tensor<4x?x32x8x16xf16>