Wip/gemma3 rebase

compiler/plugins/input/StableHLO/Conversion/Preprocessing/Canonicalization.cpp

@@ -1176,7 +1176,7 @@ struct StableHLOCanonicalize final
 };
 
 } // namespace
-void populateCanonicalizationPatterns(MLIRContext *context,
+void populateCanonicalizationPatternsNoReorder(MLIRContext *context,
                                       RewritePatternSet *patterns,
                                       PatternBenefit benefit) {
   patterns->add<
@@ -1197,6 +1197,12 @@ void populateCanonicalizationPatterns(MLIRContext *context,
       ReshapeOpCanon, MergeConsecutiveReshapes, TransposeIsReshape,
       // Types.
       ZeroExtentTensorCanon>(context, benefit);
+}
+
+void populateCanonicalizationPatterns(MLIRContext *context,
+                                      RewritePatternSet *patterns,
+                                      PatternBenefit benefit) {
+  populateCanonicalizationPatternsNoReorder(context, patterns, benefit);
   patterns->add<ReorderElementwiseAndShapeOp>(context);
 }
 } // namespace mlir::iree_compiler::stablehlo

compiler/plugins/input/StableHLO/Conversion/Preprocessing/Rewriters.h

@@ -20,6 +20,11 @@ namespace mlir::iree_compiler::stablehlo {
 void populateCanonicalizationPatterns(MLIRContext *context,
                                       RewritePatternSet *patterns,
                                       PatternBenefit benefit = 1);
+/// Collection of canonicalization patterns for StableHLO
+/// without ReorderElementwiseAndShapeOp.
+void populateCanonicalizationPatternsNoReorder(MLIRContext *context,
+                                      RewritePatternSet *patterns,
+                                      PatternBenefit benefit = 1);
 
 /// Collection of rewrite patterns for lowering of StableHLO dot general
 /// operations.

compiler/plugins/input/StableHLO/Conversion/Preprocessing/StableHLOToStableHLO.cpp

@@ -564,6 +564,16 @@ struct ScatterImplicitBatch final
     auto indices = cast<Value>(op.getScatterIndices());
     auto indicesTy = dyn_cast<RankedTensorType>(indices.getType());
 
+    auto input = op.getInputs().front();
+    auto update = op.getUpdates().front();
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    auto updateTy = dyn_cast<RankedTensorType>(update.getType());
+
+    if (inputTy.getRank() > 1 && indicesTy.getRank() == 1 &&
+        updateTy.getRank() == inputTy.getRank()) {
+      return rewriter.notifyMatchFailure(op, "special case handling.");
+    }
+
     // Check whether indices has no batch dimension.
     if (!indicesTy)
       return failure();
@@ -853,8 +863,15 @@ struct ScatterMaterializeInsertedDim final
                                 PatternRewriter &rewriter) const override {
     auto indices = op.getScatterIndices();
     auto operand = op.getInputs().front();
+    auto update = op.getUpdates().front();
     auto indicesTy = cast<ShapedType>(indices.getType());
     auto operandTy = cast<ShapedType>(operand.getType());
+    auto updateTy = cast<ShapedType>(update.getType());
+
+    if (operandTy.getRank() == updateTy.getRank() &&
+        operandTy.getRank() == indicesTy.getRank()) {
+      return rewriter.notifyMatchFailure(op, "special case handling.");
+    }
 
     if (!operandTy.hasRank() || !indicesTy.hasRank()) {
       return rewriter.notifyMatchFailure(op, "operand/indices have no rank");

compiler/plugins/input/StableHLO/Conversion/StableHLOCustomCalls.cpp

@@ -7,6 +7,7 @@
 // Implements logic for lowering CHLO ops to StableHLO and Shape dialect ops,
 // taking care of CHLO's broadcasting semantics
 
+#include <llvm/Support/Casting.h>
 #include "compiler/plugins/input/StableHLO/Conversion/Passes.h"
 #include "compiler/plugins/input/StableHLO/Conversion/Preprocessing/Rewriters.h"
 #include "compiler/plugins/input/StableHLO/Conversion/Rewriters.h"
@@ -15,6 +16,7 @@
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/TypeUtilities.h"
@@ -238,6 +240,50 @@ struct ShapeAssertionDrop final
   }
 };
 
+struct NoOpShardingDrop final
+    : OpRewritePattern<mlir::stablehlo::CustomCallOp> {
+  using Base::Base;
+  using OpAdaptor = mlir::stablehlo::CustomCallOp::Adaptor;
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::CustomCallOp op,
+                                PatternRewriter &rewriter) const final {
+    if (op.getCallTargetName() != "Sharding") {
+      return rewriter.notifyMatchFailure(op, "not ShardingCustomCall");
+    }
+    unsigned numOperands = op.getNumOperands();
+    if (numOperands != 1) {
+      return rewriter.notifyMatchFailure(
+          op, "ShardingCustomCall with more than one operand");
+    }
+
+    // Get root of the operations
+    mlir::Operation *root = op.getOperation();
+    while (root->getParentOp() != nullptr) {
+      root = root->getParentOp();
+    }
+    auto rootAttrs =
+        llvm::dyn_cast<mlir::DictionaryAttr>(root->getAttrDictionary());
+    auto numPartitions = llvm::dyn_cast_if_present<mlir::IntegerAttr>(
+        rootAttrs.get("mhlo.num_partitions"));
+    auto numReplicas = llvm::dyn_cast_if_present<mlir::IntegerAttr>(
+        rootAttrs.get("mhlo.num_replicas"));
+    if (numPartitions == nullptr || numReplicas == nullptr) {
+      return rewriter.notifyMatchFailure(
+          op,
+          "ShardingCustomCall: Number of partitions or replicas not specified");
+    }
+
+    if (numPartitions.getInt() == 1 && numReplicas.getInt() == 1) {
+      // Remove sharding as number of partitions and replicas is 1
+      rewriter.replaceOp(op, op.getOperands()[0]);
+      return success();
+    }
+    // There is more than one partition or replica
+    return rewriter.notifyMatchFailure(
+        op, "Sharding with multiple partitions or replicas is not supported");
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // Pass Definition.
 //===----------------------------------------------------------------------===//
@@ -254,7 +300,8 @@ struct LegalizeStableHLOCustomCalls final
     MLIRContext *ctx = f.getContext();
 
     RewritePatternSet patterns(ctx);
-    patterns.add<HouseholderReflectorRewriter, ShapeAssertionDrop>(ctx);
+    patterns.add<HouseholderReflectorRewriter, ShapeAssertionDrop,
+                 NoOpShardingDrop>(ctx);
     if (failed(applyPatternsGreedily(f, std::move(patterns)))) {
       signalPassFailure();
     }

compiler/plugins/input/StableHLO/Conversion/StableHLOToIREEInputDialects.cpp

@@ -511,7 +511,7 @@ struct ConvertStableHloToIreeInputDialects final
 
     // Run stablehlo canonicalization patterns with a high benefit to avoid some
     // expensive expansions.
-    populateCanonicalizationPatterns(context, &patterns, /*benefit=*/1024);
+    populateCanonicalizationPatternsNoReorder(context, &patterns, /*benefit=*/1024);
 
     // Run custom patterns with a high benefit to override stablehlo patterns.
     patterns.add<ConcatenateOpConversion, FftOpConversion,

compiler/plugins/input/StableHLO/Conversion/StableHLOToLinalgExt.cpp

@@ -247,8 +247,6 @@ struct ScatterOpConversion final
   LogicalResult
   matchAndRewrite(mlir::stablehlo::ScatterOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    if (!hasCanonicalDimensionNumbers(op))
-      return failure();
     if (llvm::size(op.getInputs()) != 1)
       return op.emitError("NYI variadic operands scatter");
     if (llvm::size(op.getUpdates()) != 1)
@@ -260,7 +258,149 @@ struct ScatterOpConversion final
     Value indices = adaptor.getScatterIndices();
     Value updates = adaptor.getUpdates().front();
 
-    auto originalType = dyn_cast<ShapedType>(original.getType());
+    Location loc = op.getLoc();
+    auto originalType = dyn_cast<RankedTensorType>(original.getType());
+    auto indicesType = dyn_cast<RankedTensorType>(indices.getType());
+    auto updateType = dyn_cast<RankedTensorType>(updates.getType());
+
+    auto dimNumbers = op.getScatterDimensionNumbers();
+    auto updateDims = dimNumbers.getUpdateWindowDims();
+
+    if (originalType.getRank() == 2 && //
+        originalType.getShape().front() == 1 &&
+        updateType.getShape().front() == 1 &&
+        indicesType.getShape().front() == 1 && //
+        updateDims.size() == 1 &&              //
+        updateDims.front() == 1 &&
+        dimNumbers.getInsertedWindowDims().size() == 1 &&
+        dimNumbers.getInsertedWindowDims().front() == 1 &&
+        dimNumbers.getInputBatchingDims().empty() &&
+        dimNumbers.getScatterIndicesBatchingDims().empty() &&
+        dimNumbers.getScatterDimsToOperandDims().size() == 1 &&
+        dimNumbers.getScatterDimsToOperandDims().front() == 1) {
+
+      auto newIndices = mlir::tensor::CollapseShapeOp::create(
+          rewriter, loc, op.getScatterIndices(),
+          mlir::ArrayRef<mlir::ReassociationIndices>{{0, 1}});
+
+      Value zero = mlir::arith::ConstantIndexOp::create(rewriter, loc, 0);
+      Value newIndices0 =
+          mlir::tensor::ExtractOp::create(rewriter, loc, newIndices, zero);
+
+      auto newUpdates = mlir::tensor::CollapseShapeOp::create(
+          rewriter, loc, updates,
+          mlir::ArrayRef<mlir::ReassociationIndices>{{0, 1}});
+
+      SmallVector<OpFoldResult, 4> sizes;
+      for (int64_t size : updateType.getShape()) {
+        sizes.push_back(rewriter.getIndexAttr(size));
+      }
+
+      newIndices0 = mlir::arith::IndexCastOp::create(
+          rewriter, loc, rewriter.getIndexType(), newIndices0);
+      SmallVector<OpFoldResult, 4> offsets = {zero, newIndices0};
+
+      int64_t rank = originalType.getRank();
+      SmallVector<OpFoldResult, 4> strides(rank, rewriter.getI64IntegerAttr(1));
+      rewriter.replaceOpWithNewOp<mlir::tensor::InsertSliceOp>(
+          op, newUpdates, original, offsets, sizes, strides);
+
+      return success();
+    }
+    if (originalType.getRank() == 2 && //
+        originalType.getShape().front() == 1 &&
+        updateType.getShape().front() == 2 &&
+        indicesType.getShape().front() == 2 && //
+        updateDims.size() == 1 &&              //
+        updateDims.front() == 1 &&
+        dimNumbers.getInsertedWindowDims().size() == 1 &&
+        dimNumbers.getInsertedWindowDims().front() == 1 &&
+        dimNumbers.getInputBatchingDims().empty() &&
+        dimNumbers.getScatterIndicesBatchingDims().empty() &&
+        dimNumbers.getScatterDimsToOperandDims().size() == 1 &&
+        dimNumbers.getScatterDimsToOperandDims().front() == 1) {
+
+      Value zero = mlir::arith::ConstantIndexOp::create(rewriter, loc, 0);
+      Value one = mlir::arith::ConstantIndexOp::create(rewriter, loc, 1);
+
+      auto newIndices = mlir::tensor::CollapseShapeOp::create(
+          rewriter, loc, op.getScatterIndices(),
+          mlir::ArrayRef<mlir::ReassociationIndices>{{0, 1}});
+
+      Value newIndices0 =
+          mlir::tensor::ExtractOp::create(rewriter, loc, newIndices, zero);
+      newIndices0 = mlir::arith::IndexCastOp::create(
+          rewriter, loc, rewriter.getIndexType(), newIndices0);
+      Value newIndices1 =
+          mlir::tensor::ExtractOp::create(rewriter, loc, newIndices, one);
+      newIndices1 = mlir::arith::IndexCastOp::create(
+          rewriter, loc, rewriter.getIndexType(), newIndices1);
+
+      SmallVector<int64_t> vShape = {1, 1};
+      SmallVector<OpFoldResult> vSizes;
+      for (auto v : vShape) {
+        vSizes.push_back(b.getIndexAttr(v));
+      }
+
+      auto sliceTy = RankedTensorType::get(vShape, updateType.getElementType());
+      SmallVector<OpFoldResult> vOffsets0(updateType.getRank(),
+                                          b.getIndexAttr(0));
+      SmallVector<OpFoldResult> vStrides(updateType.getRank(),
+                                         b.getIndexAttr(1));
+      auto newUpdates0 = mlir::tensor::ExtractSliceOp::create(
+          rewriter, loc, sliceTy, updates, vOffsets0, vSizes, vStrides);
+      SmallVector<OpFoldResult> vOffsets1(updateType.getRank(),
+                                          b.getIndexAttr(0));
+      vOffsets1[0] = b.getIndexAttr(1);
+      auto newUpdates1 = mlir::tensor::ExtractSliceOp::create(
+          rewriter, loc, sliceTy, updates, vOffsets1, vSizes, vStrides);
+
+      SmallVector<OpFoldResult, 4> offsets0 = {zero, newIndices0};
+      SmallVector<OpFoldResult, 4> offsets1 = {zero, newIndices1};
+      for (int i = 0; i < originalType.getRank() - 2; i++) {
+        offsets0.push_back(zero);
+        offsets1.push_back(zero);
+      }
+
+      int64_t rank = originalType.getRank();
+      SmallVector<OpFoldResult, 4> sizes = {b.getIndexAttr(1),
+                                            b.getIndexAttr(1)};
+      SmallVector<OpFoldResult, 4> strides(rank, rewriter.getI64IntegerAttr(1));
+      Value slice = mlir::tensor::InsertSliceOp::create(
+          b, newUpdates0, original, offsets0, sizes, strides);
+      rewriter.replaceOpWithNewOp<mlir::tensor::InsertSliceOp>(
+          op, newUpdates1, slice, offsets1, sizes, strides);
+
+      return success();
+    }
+
+    if (originalType.getRank() > 1 && indicesType.getRank() == 1 &&
+        updateType.getRank() == originalType.getRank()) {
+      SmallVector<OpFoldResult, 4> sizes;
+      for (int64_t size : updateType.getShape()) {
+        sizes.push_back(rewriter.getIndexAttr(size));
+      }
+
+      Value zero = mlir::arith::ConstantIndexOp::create(rewriter, loc, 0);
+      Value extracted =
+          mlir::tensor::ExtractOp::create(rewriter, loc, indices, zero);
+      extracted = mlir::arith::IndexCastOp::create(
+          rewriter, loc, rewriter.getIndexType(), extracted);
+      SmallVector<OpFoldResult, 4> offsets = {zero, extracted};
+      for (int i = 0; i < originalType.getRank() - 2; i++) {
+        offsets.push_back(zero);
+      }
+
+      int64_t rank = originalType.getRank();
+      SmallVector<OpFoldResult, 4> strides(rank, rewriter.getI64IntegerAttr(1));
+      rewriter.replaceOpWithNewOp<mlir::tensor::InsertSliceOp>(
+          op, updates, original, offsets, sizes, strides);
+
+      return success();
+    }
+
+    if (!hasCanonicalDimensionNumbers(op))
+      return failure();
 
     llvm::SmallVector<int64_t> scatterDimMap;
     for (auto dim :

compiler/plugins/input/StableHLO/Conversion/test/stablehlo_custom_calls.mlir

@@ -10,3 +10,13 @@ func.func public @householder(%arg0: tensor<4x3xf32>, %arg1: tensor<2xf32>) -> (
   %0 = stablehlo.custom_call @ProductOfElementaryHouseholderReflectors(%arg0, %arg1) : (tensor<4x3xf32>, tensor<2xf32>) -> tensor<4x3xf32>
   return %0 : tensor<4x3xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @noop_sharding_custom_call
+func.func public @noop_sharding_custom_call(%arg0: tensor<2xui32>) {
+  %result = "stablehlo.custom_call"(%arg0) <{call_target_name = "Sharding"}> {mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@empty_mesh, [{}]>]>"}, mhlo.sharding = "{replicated}"} : (tensor<2xui32>) -> tensor<2xui32>
+  // CHECK-NOT: stablehlo.custom_call
+  // CHECK-NOT: sharding
+  return
+}

compiler/plugins/input/StableHLO/Conversion/test/stablehlo_to_iree_input_dialects.mlir

@@ -140,3 +140,20 @@ func.func @select_conversion(%arg0: tensor<i1>, %arg1: tensor<?xui16>, %arg2: te
   %0 = arith.select %extracted, %arg1, %arg2 : tensor<?xui16>
   return %0 : tensor<?xui16>
 }
+
+// -----
+
+// Tests whether the pass fails on ReorderElementwiseAndShapeOp pattern.
+
+// CHECK:       func.func @reorder_elementwise_and_shape
+// CHECK-SAME:    %[[ARG0:[^:]+]]: tensor<2048xi1>
+// CHECK-SAME:    %[[ARG1:[^:]+]]: tensor<1x2048xi32>
+func.func @reorder_elementwise_and_shape(%arg0: tensor<2048xi1>, %arg1: tensor<1x2048xi32>) -> tensor<1x2048xi32> {
+  // CHECK: tensor.expand_shape
+  %0 = stablehlo.reshape %arg0 : (tensor<2048xi1>) -> tensor<1x2048xi1>
+  // CHECK: linalg.generic
+  %1 = stablehlo.convert %0 : (tensor<1x2048xi1>) -> tensor<1x2048xi32>
+  // CHECK: linalg.generic
+  %2 = stablehlo.subtract %arg1, %1 : tensor<1x2048xi32>
+  return %2 : tensor<1x2048xi32>
+}

compiler/src/iree/compiler/Dialect/Stream/Analysis/Affinity.cpp

@@ -538,6 +538,9 @@ TraversalResult ValueConsumerAffinityPVS::updateFromUse(Value value,
           auto &valueUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
               *this, value, DFX::Resolution::REQUIRED);
           newState ^= valueUsage.getState();
+          if (operand.getOperandNumber() >= whileOp->getResults().size()) {
+            return TraversalResult::INCOMPLETE;
+          }
           auto &parentUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
               *this,
               Position::forValue(

compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.cpp

@@ -1210,8 +1210,10 @@ TraversalResult Explorer::walkTransitiveUses(Value value, UseWalkFn fn,
       if (ownerOp->hasTrait<OpTrait::ReturnLike>() &&
           !isa<CallableOpInterface>(ownerOp->getParentOp())) {
         auto parent = ownerOp->getParentOp();
-        auto result = parent->getResult(use.getOperandNumber());
-        worklist.insert(result);
+        if (use.getOperandNumber() < parent->getResults().size()) {
+          auto result = parent->getResult(use.getOperandNumber());
+          worklist.insert(result);
+        }
       }
 
       // Step across global stores and into all of the loads across the program.

compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.h

@@ -367,7 +367,14 @@ static inline int32_t getRoundedElementByteWidth(Type type) {
     return vectorType.getNumElements() *
            getRoundedElementByteWidth(vectorType.getElementType());
   }
-  unsigned bitsUnaligned = type.getIntOrFloatBitWidth();
+
+  unsigned bitsUnaligned;
+  if (type.isIndex()) {
+    bitsUnaligned = dyn_cast<IndexType>(type).kInternalStorageBitWidth;
+  } else {
+    bitsUnaligned = type.getIntOrFloatBitWidth();
+  }
+
   assert(bitsUnaligned > 0 && "0-width types unsupported");
   // Round up to 8-bit aligned bytes.
   unsigned byteAligned = (bitsUnaligned + 8 - 1) / 8;