[Mlir-commits] [mlir] [mlir][linalg] Scaled contraction op (PR #203958)

[Adam Siemieniuk]

================
@@ -6842,6 +6868,296 @@ Speculation::Speculatability BatchReduceMatmulOp::getSpeculatability() {
   return getGenericSpeculatabilityImpl(cast<LinalgOp>(getOperation()));
 }
 
+//===----------------------------------------------------------------------===//
+// ScaledContractOp
+//===----------------------------------------------------------------------===//
+
+SmallVector<utils::IteratorType> ScaledContractOp::getIteratorTypesArray() {
+  AffineMap outAffineMap = getIndexingMapsArray().pop_back_val();
+  // Infer iterator types based on the output.
+  SmallVector<bool> dimsInOutput(outAffineMap.getNumDims(), false);
+  for (auto result : outAffineMap.getResults()) {
+    auto dimExpr = dyn_cast<AffineDimExpr>(result);
+    assert(dimExpr && "affine_map is a projected permutation");
+    dimsInOutput[dimExpr.getPosition()] = true;
+  }
+
+  SmallVector<utils::IteratorType> iteratorTypes;
+  for (auto dimOccursInOutput : dimsInOutput)
+    iteratorTypes.push_back(dimOccursInOutput ? utils::IteratorType::parallel
+                                              : utils::IteratorType::reduction);
+
+  return iteratorTypes;
+}
+
+unsigned ScaledContractOp::getNumRegionArgs() { return 5; }
+
+/// Implement block region builder, which is called by 'fillStructuredOpRegion'.
+void ScaledContractOp::regionBuilder(
+    ImplicitLocOpBuilder &b, Block &block, ArrayRef<NamedAttribute> attrs,
+    function_ref<InFlightDiagnostic()> emitError) {
+  if (emitError && block.getNumArguments() != 5) {
+    emitError() << "ScaledContractOp regionBuilder expects 5 args, got "
+                << block.getNumArguments();
+    return;
+  }
+  assert(block.getNumArguments() == 5 &&
+         "ScaledContractOp regionBuilder expects 5 args");
+  RegionBuilderHelper helper(b, block);
+
+  TypeFn castSignedness = TypeFn::cast_signed;
+  auto castIter = llvm::find_if(attrs, [&](const NamedAttribute &attr) {
+    return attr.getName() == "cast";
+  });
+  if (castIter != attrs.end()) {
+    if (auto attr = llvm::dyn_cast<TypeFnAttr>(castIter->getValue()))
+      castSignedness = attr.getValue();
+  }
+
+  // TODO: Support fields with operators besides mult & add.
+  Type outType = block.getArgument(4).getType();
+
+  // Build input data value scaling.
+  // Uses specialized arith ops when possible.
+  // Otherwise, constructs computation manually.
+  auto buildScaledValue = [&](Value data, Value scale) -> Value {
+    auto dataFloatTy = dyn_cast<FloatType>(data.getType());
+    auto outFloatTy = dyn_cast<FloatType>(outType);
+    if (dataFloatTy && dyn_cast<FloatType>(scale.getType()) && outFloatTy) {
+      unsigned dataWidth = dataFloatTy.getWidth();
+      unsigned outWidth = outFloatTy.getWidth();
+      if (dataWidth < outWidth)
+        return arith::ScalingExtFOp::create(b, outType, data, scale,
+                                            /*fastmath=*/nullptr);
+      if (dataWidth > outWidth)
----------------
adam-smnk wrote:

> A more fundamental question is why you allow this at all (input data precision higher than accumulation precision)?

I started with the most open design in line with rest of contraction ops that allows to generally freely mix precisions.
That being said, I see no real use case where that'd be needed.

I'll restrict accumulator type to be equal or higher precision than inputs.

https://github.com/llvm/llvm-project/pull/203958