testing PR700

lib/Conversion/OnnxToHip/NormConversion.cpp

@@ -51,6 +51,51 @@ inline mlir::Value getOptionalResult(mlir::Operation *op, unsigned idx) {
   return v;
 }
 
+/// Reject a norm whose `axis` does not select exactly the trailing block that
+/// the scale operand covers.
+///
+/// The runtime wrappers flatten the input to `[num_rows, hidden_dim]` with
+/// `hidden_dim` taken from the scale's element count and `num_rows` recovered
+/// by dividing the input's total element count. That is only the requested
+/// reduction when normalization spans `dims[axis:]` and those dims multiply to
+/// the scale's size; the wrappers cannot check it themselves because the ABI
+/// carries element counts rather than shapes. Rank is known here, so check it
+/// here -- an unnoticed mismatch normalizes over the wrong axis and returns
+/// plausible-looking numbers.
+///
+/// Dynamic trailing extents are not checkable and are accepted; in practice the
+/// normalized axes are the static feature dims.
+mlir::LogicalResult verifyNormAxisMatchesScale(mlir::Operation *op,
+                                               mlir::Value input,
+                                               mlir::Value scale,
+                                               int64_t axis) {
+  auto inputType = mlir::dyn_cast<mlir::RankedTensorType>(input.getType());
+  auto scaleType = mlir::dyn_cast<mlir::RankedTensorType>(scale.getType());
+  if (!inputType || !scaleType || !scaleType.hasStaticShape())
+    return mlir::success();
+
+  int64_t rank = inputType.getRank();
+  int64_t normAxis = axis < 0 ? axis + rank : axis;
+  if (normAxis < 0 || normAxis >= rank)
+    return op->emitError("norm axis ")
+           << axis << " is out of range for a rank-" << rank << " input";
+
+  int64_t normalizedExtent = 1;
+  for (int64_t d = normAxis; d < rank; ++d) {
+    if (inputType.isDynamicDim(d))
+      return mlir::success();
+    normalizedExtent *= inputType.getDimSize(d);
+  }
+  if (normalizedExtent != scaleType.getNumElements())
+    return op->emitError("norm axis ")
+           << axis << " spans " << normalizedExtent
+           << " element(s) of the rank-" << rank << " input, but scale has "
+           << scaleType.getNumElements()
+           << "; the runtime flattens to [rows, scale_num_elements] and would "
+              "normalize over the wrong axis";
+  return mlir::success();
+}
+
 /// onnx.Custom(SimplifiedLayerNormalization) -> hip.rms_norm
 struct SimplifiedLayerNormToHip : public mlir::RewritePattern {
   SimplifiedLayerNormToHip(mlir::MLIRContext *ctx)
@@ -98,6 +143,10 @@ mlir::LogicalResult SimplifiedLayerNormToHip::matchAndRewrite(
   if (!stashTypeAttr)
     return rewriter.notifyMatchFailure(op, "missing stash_type attribute");
 
+  if (mlir::failed(
+          verifyNormAxisMatchesScale(op, input, scale, axisAttr.getSInt())))
+    return mlir::failure();
+
   // Convert axis to i64
   auto axisI64Attr = rewriter.getI64IntegerAttr(axisAttr.getSInt());
   auto stashTypeI64Attr = rewriter.getI64IntegerAttr(stashTypeAttr.getSInt());
@@ -153,6 +202,9 @@ RMSNormalizationToHip::matchAndRewrite(mlir::Operation *op,
   if (auto axisAttr = op->getAttrOfType<mlir::IntegerAttr>("axis"))
     axis = axisAttr.getSInt();
 
+  if (mlir::failed(verifyNormAxisMatchesScale(op, input, scale, axis)))
+    return mlir::failure();
+
   llvm::APFloat epsValue(9.99999974E-6f);
   if (auto epsilonAttr = op->getAttrOfType<mlir::FloatAttr>("epsilon"))
     epsValue = epsilonAttr.getValue();

lib/Conversion/OnnxToHip/ReshapeConversion.cpp

@@ -51,23 +51,106 @@ validateSqueezeUnsqueezeOp(mlir::Operation *op, mlir::PatternRewriter &rewriter,
   return mlir::success();
 }
 
+/// Count the dynamic output dims a reassociation group covers.
+static int64_t countDynOutDims(mlir::RankedTensorType outputType,
+                               const mlir::ReassociationIndices &group) {
+  return llvm::count_if(
+      group, [&](int64_t idx) { return outputType.isDynamicDim(idx); });
+}
+
+/// Resolve a Reshape's `shape` operand into one extent per output dim.
+///
+/// Only a host-visible shape vector is honoured, i.e. the
+/// `tensor.from_elements` that ReshapeShapeFold leaves behind for the
+/// `Reshape(_, Shape(x))` idiom. A shape tensor that is still device-resident
+/// yields nullopt instead of paying for a readback; the caller then falls back
+/// to `tensor.reshape`, which consumes the runtime shape vector directly.
+static std::optional<llvm::SmallVector<mlir::OpFoldResult>>
+resolveShapeOperandExtents(mlir::PatternRewriter &rewriter, mlir::Location loc,
+                           mlir::Value shapeOperand, int64_t outputRank) {
+  if (!shapeOperand)
+    return std::nullopt;
+  auto fromElements =
+      shapeOperand.getDefiningOp<mlir::tensor::FromElementsOp>();
+  if (!fromElements ||
+      static_cast<int64_t>(fromElements.getElements().size()) != outputRank)
+    return std::nullopt;
+
+  // Resolve every element before building anything, so a late unresolvable
+  // entry cannot leave a half-materialized cast behind in the IR.
+  llvm::SmallVector<mlir::OpFoldResult> extents;
+  llvm::SmallVector<mlir::Value> needsCast(outputRank);
+  extents.reserve(outputRank);
+  for (auto [i, element] : llvm::enumerate(fromElements.getElements())) {
+    // ReshapeShapeFold emits `arith.index_cast %tensor.dim`. Reusing the
+    // pre-cast index keeps the extent on the same SSA value the rest of the
+    // conversion derives dims from, instead of a round trip through i64.
+    if (auto cast = element.getDefiningOp<mlir::arith::IndexCastOp>();
+        cast && mlir::isa<mlir::IndexType>(cast.getIn().getType())) {
+      extents.push_back(cast.getIn());
+      continue;
+    }
+    if (auto constant = element.getDefiningOp<mlir::arith::ConstantOp>()) {
+      auto intAttr = mlir::dyn_cast<mlir::IntegerAttr>(constant.getValue());
+      // ONNX gives 0 ("keep the input dim") and -1 ("infer") meanings that
+      // this helper does not resolve; neither is usable as an extent.
+      if (!intAttr || intAttr.getInt() <= 0)
+        return std::nullopt;
+      extents.push_back(rewriter.getIndexAttr(intAttr.getInt()));
+      continue;
+    }
+    needsCast[i] = element;
+    extents.push_back(mlir::OpFoldResult{});
+  }
+
+  for (int64_t i : llvm::seq<int64_t>(outputRank)) {
+    if (!needsCast[i])
+      continue;
+    mlir::Value asIndex = mlir::arith::IndexCastOp::create(
+        rewriter, loc, rewriter.getIndexType(), needsCast[i]);
+    extents[i] = asIndex;
+  }
+  return extents;
+}
+
 /// Build output shape for expand_shape operations.
 /// Used by both Reshape and Unsqueeze when expanding dimensions.
 ///
 /// For static dimensions: use compile-time size from outputType.
 /// For dynamic dimensions: extract from input via DimOp, dividing out any
 /// static dimensions in the same reassociation group.
-llvm::SmallVector<mlir::OpFoldResult> buildExpandShapeOutputShape(
-    mlir::PatternRewriter &rewriter, mlir::Location loc, mlir::Value data,
-    mlir::RankedTensorType outputType,
-    llvm::ArrayRef<mlir::ReassociationIndices> reassoc) {
+///
+/// That derivation only holds while a group covers at most ONE dynamic output
+/// dim. When a group splits one dynamic source dim into several dynamic output
+/// dims, the source extent is their PRODUCT and says nothing about the split,
+/// so those extents have to come from \p shapeOperand (the Reshape's `shape`
+/// input). Returns nullopt when that is needed but unreadable, letting the
+/// caller fall back to `tensor.reshape` instead of emitting an expand_shape
+/// that claims each dynamic output dim is the whole product -- which is how
+/// `[bs*ss, 2816] -> [bs, ss, 2816]` became `[ss, ss, 2816]` and dispatched
+/// every Gemma-4 layer's input norm over ss^2 rows.
+std::optional<llvm::SmallVector<mlir::OpFoldResult>>
+buildExpandShapeOutputShape(mlir::PatternRewriter &rewriter, mlir::Location loc,
+                            mlir::Value data, mlir::RankedTensorType outputType,
+                            llvm::ArrayRef<mlir::ReassociationIndices> reassoc,
+                            mlir::Value shapeOperand = {}) {
   int64_t outputRank = outputType.getRank();
 
   llvm::SmallVector<int64_t> outDimToInDim(outputRank, -1);
   for (auto [g, group] : llvm::enumerate(reassoc))
     for (int64_t idx : group)
       outDimToInDim[idx] = g;
 
+  std::optional<llvm::SmallVector<mlir::OpFoldResult>> shapeExtents;
+  if (llvm::any_of(reassoc, [&](const mlir::ReassociationIndices &group) {
+        return countDynOutDims(outputType, group) > 1;
+      })) {
+    shapeExtents =
+        resolveShapeOperandExtents(rewriter, loc, shapeOperand, outputRank);
+    if (!shapeExtents)
+      return std::nullopt;
+  }
+
   llvm::SmallVector<mlir::OpFoldResult> outputShape;
   for (int64_t i : llvm::seq<int64_t>(outputRank)) {
     if (!outputType.isDynamicDim(i)) {
@@ -78,6 +161,11 @@ llvm::SmallVector<mlir::OpFoldResult> buildExpandShapeOutputShape(
     int64_t srcDim = outDimToInDim[i];
     const auto &group = reassoc[srcDim];
 
+    if (countDynOutDims(outputType, group) > 1) {
+      outputShape.push_back((*shapeExtents)[i]);
+      continue;
+    }
+
     int64_t staticProduct = 1;
     for (int64_t idx : group)
       if (!outputType.isDynamicDim(idx))
@@ -194,17 +282,27 @@ struct ReshapeToStdTensor : public mlir::RewritePattern {
       if (auto reassocOpt =
               mlir::getReassociationIndicesForReshape(inputType, outputType)) {
         if (outputRank > inputRank) {
-          auto outputShape = buildExpandShapeOutputShape(
-              rewriter, loc, data, outputType, *reassocOpt);
-          auto expandOp = mlir::tensor::ExpandShapeOp::create(
-              rewriter, loc, outputType, data, *reassocOpt, outputShape);
-          rewriter.replaceOp(op, expandOp.getResult());
+          // The `shape` operand is the only place a multi-dynamic split's
+          // per-dim extents exist; ReshapeShapeFold has already folded
+          // `Shape(x)` into a host-visible tensor.from_elements for it.
+          mlir::Value shapeOperand =
+              op->getNumOperands() >= 2 ? op->getOperand(1) : mlir::Value();
+          if (auto outputShape = buildExpandShapeOutputShape(
+                  rewriter, loc, data, outputType, *reassocOpt, shapeOperand)) {
+            auto expandOp = mlir::tensor::ExpandShapeOp::create(
+                rewriter, loc, outputType, data, *reassocOpt, *outputShape);
+            rewriter.replaceOp(op, expandOp.getResult());
+            return mlir::success();
+          }
+          // Multi-dynamic split with an unreadable shape operand: fall through
+          // to the tensor.reshape fallback, which takes the runtime shape
+          // vector verbatim.
         } else {
           auto collapseOp = mlir::tensor::CollapseShapeOp::create(
               rewriter, loc, outputType, data, *reassocOpt);
           rewriter.replaceOp(op, collapseOp.getResult());
+          return mlir::success();
         }
-        return mlir::success();
       }
       // Fall through to tensor.reshape fallback (no structured reassoc).
     }
@@ -349,11 +447,15 @@ struct ReshapeToStdTensor : public mlir::RewritePattern {
         // combine direction here. (PoolAllocs's hoistable whitelist must
         // include arith.divui for the resulting dim arithmetic to survive
         // pool-base hoisting.)
+        // Each group here pairs one dynamic extent with the static factor K,
+        // so no group is multi-dynamic and the shape operand is never needed.
         auto intOutShape = buildExpandShapeOutputShape(rewriter, loc, data,
                                                        intType, expandReassoc);
+        if (!intOutShape)
+          break; // fall through to tensor.reshape fallback
 
         auto expanded = mlir::tensor::ExpandShapeOp::create(
-            rewriter, loc, intType, data, expandReassoc, intOutShape);
+            rewriter, loc, intType, data, expandReassoc, *intOutShape);
 
         // (e) Collapse: the OUTPUT dim that absorbs the factor pair maps to the
         // two intermediate dims; everything else is identity.
@@ -578,10 +680,15 @@ struct UnsqueezeToStdTensor : public mlir::RewritePattern {
           op, "cannot compute unsqueeze reassociation");
 
     mlir::Location loc = op->getLoc();
+    // Unsqueeze only inserts unit dims, so every group keeps at most one
+    // dynamic dim and no shape operand is needed.
     auto outputShape = buildExpandShapeOutputShape(rewriter, loc, data,
                                                    outputType, *reassocOpt);
+    if (!outputShape)
+      return rewriter.notifyMatchFailure(
+          op, "unsqueeze: multi-dynamic reassociation group");
     auto expandOp = mlir::tensor::ExpandShapeOp::create(
-        rewriter, loc, outputType, data, *reassocOpt, outputShape);
+        rewriter, loc, outputType, data, *reassocOpt, *outputShape);
     rewriter.replaceOp(op, expandOp.getResult());
     return mlir::success();
   }

lib/Conversion/OnnxToHip/ReshapeShapeFold.cpp

@@ -50,8 +50,9 @@
 //
 // This fold rewrites the shape operand BEFORE ConvertOnnxToHip so the
 // resulting `tensor.from_elements` is exactly what ReshapeConversion's
-// existing multi-dyn-per-group branch picks up.  Bug fix is structural
-// at the operand level; no change to ReshapeConversion required.
+// multi-dyn-per-group branch picks up.  That branch is the other half of
+// the fix and lives in `buildExpandShapeOutputShape`; without it this fold
+// has no effect, because the conversion otherwise never reads operand 1.
 //
 // Implementation notes
 // --------------------

lib/Runtime/real/simplified_layer_norm.cpp

@@ -212,6 +212,25 @@ int wrap_miopenT5LayerNormForward(RuntimeState *state, int op_state_slot,
     return -1;
   }
 
+  // The [num_rows, hidden_dim] flattening below is only meaningful while the
+  // caller's element counts agree; matches the checks in
+  // wrap_layer_normalization. The axis-vs-scale agreement cannot be rechecked
+  // here (the ABI carries element counts, not shapes) and is enforced at
+  // conversion time by verifyNormAxisMatchesScale.
+  if (scale_num_elements <= 0) {
+    fprintf(stderr,
+            "[REAL] wrap_miopenT5LayerNormForward: scale_num_elements=%lld\n",
+            (long long)scale_num_elements);
+    return -1;
+  }
+  if (input_num_elements % scale_num_elements != 0) {
+    fprintf(stderr,
+            "[REAL] wrap_miopenT5LayerNormForward: input_num(%lld) not "
+            "divisible by scale_num(%lld)\n",
+            (long long)input_num_elements, (long long)scale_num_elements);
+    return -1;
+  }
+
   int64_t hidden_dim = scale_num_elements;
   int64_t num_rows = input_num_elements / hidden_dim;
 

lib/Runtime/real/skip_simplified_layer_norm.cpp

@@ -240,6 +240,22 @@ int wrap_skip_simplified_layer_norm(RuntimeState *state, int op_state_slot,
     return -1;
   }
 
+  // Same flattening precondition as wrap_layer_normalization: without it a
+  // caller whose element counts disagree silently gets a truncated row count.
+  if (gamma_num_elements <= 0) {
+    fprintf(stderr,
+            "[REAL] wrap_skip_simplified_layer_norm: gamma_num_elements=%lld\n",
+            (long long)gamma_num_elements);
+    return -1;
+  }
+  if (input_num_elements % gamma_num_elements != 0) {
+    fprintf(stderr,
+            "[REAL] wrap_skip_simplified_layer_norm: input_num(%lld) not "
+            "divisible by gamma_num(%lld)\n",
+            (long long)input_num_elements, (long long)gamma_num_elements);
+    return -1;
+  }
+
   int64_t hidden_dim = gamma_num_elements;
   int64_t num_rows = input_num_elements / hidden_dim;
 

test/lit/Conversion/hip-to-llvm/test_rms_norm.mlir

@@ -113,6 +113,39 @@ func.func @rms_norm_dynamic(%ctx: !hip.context, %input: memref<?x512xf16, 1>) {
   return
 }
 
+// Rank-3 [batch, seq, hidden] with both leading extents dynamic: the shape a
+// decoder layer's input norm actually has. input_num_elements must be the
+// product of THREE distinct descriptor reads -- if dim 0 and dim 1 resolve to
+// the same value the norm is dispatched over seq^2 rows.
+// CHECK-LABEL: @rms_norm_dynamic_batch_seq_3d
+func.func @rms_norm_dynamic_batch_seq_3d(%ctx: !hip.context,
+                                         %input: memref<?x?x2816xf16, 1>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2816 = arith.constant 2816 : index
+  %dim0 = memref.dim %input, %c0 : memref<?x?x2816xf16, 1>
+  %dim1 = memref.dim %input, %c1 : memref<?x?x2816xf16, 1>
+  %scale = memref.alloc(%c2816) : memref<?xf16, 1>
+  %output = memref.alloc(%dim0, %dim1) : memref<?x?x2816xf16, 1>
+
+  // Batch and sequence come from separate descriptor slots.
+  // CHECK: llvm.extractvalue {{.*}}[3, 0]
+  // CHECK: llvm.extractvalue {{.*}}[3, 1]
+
+  // input_num_elements = dim0 * dim1 * 2816
+  // CHECK: llvm.mul {{.*}}, {{.*}} : i64
+  // CHECK: llvm.mul {{.*}}, {{.*}} : i64
+  // CHECK: llvm.mul {{.*}}, {{.*}} : i64
+
+  // CHECK: llvm.call @wrap_miopenT5LayerNormForward({{.*}}) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, f32, i64) -> i32
+  hip.rms_norm(%ctx)
+      ins(%input, %scale : memref<?x?x2816xf16, 1>, memref<?xf16, 1>)
+      outs(%output : memref<?x?x2816xf16, 1>)
+      {axis = -1 : i64, epsilon = 9.99999997e-07 : f32, stash_type = 1 : i64}
+
+  return
+}
+
 // CHECK-LABEL: @rms_norm_fully_dynamic
 func.func @rms_norm_fully_dynamic(%ctx: !hip.context, %input: memref<?x?xf16, 1>) {
   %c0 = arith.constant 0 : index