[MLIR][Shape] Fold `shape.mul`

[frgossen]

Summary

Implement constant folding for shape.mul.

Depends On D84437

[frgossen]

Rebase

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_MulOp : Shape_Op<"mul", [Commutative, NoSideEffect]> {
   let arguments = (ins Shape_SizeOrIndexType:$lhs, Shape_SizeOrIndexType:$rhs);
   let results = (outs Shape_SizeOrIndexType:$result);
 
   let assemblyFormat = [{
     $lhs `,` $rhs `:` type($lhs) `,` type($rhs) `->` type($result) attr-dict
   }];
 
   let verifier = [{ return ::verify(*this); }];
+  let hasFolder = 1;
 }
 
 def Shape_NumElementsOp : Shape_Op<"num_elements", [NoSideEffect]> {
   let summary = "Returns the number of elements for a given shape";
   let description = [{
     Returns the number of elements for a given shape which is the product of its
     dimensions.
 

mlir/lib/Dialect/Shape/IR/Shape.cpp

static LogicalResult verify(MulOp op) {
   } else {
     if (resultTy.isa<SizeType>())
       return op.emitError() << "if none of the operands can hold error values "
                                "then the result must be of type `index`";
   }
   return success();
 }
 
+OpFoldResult MulOp::fold(ArrayRef<Attribute> operands) {
+  auto lhs = operands[0].dyn_cast_or_null<IntegerAttr>();
+  if (!lhs)
+    return nullptr;
+  auto rhs = operands[1].dyn_cast_or_null<IntegerAttr>();
+  if (!rhs)
+    return nullptr;
+  APInt folded = lhs.getValue() * rhs.getValue();
+  Type indexTy = IndexType::get(getContext());
+  return IntegerAttr::get(indexTy, folded);
+}
+
 //===----------------------------------------------------------------------===//
 // ShapeOfOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ShapeOfOp::fold(ArrayRef<Attribute>) {
   auto type = getOperand().getType().dyn_cast<ShapedType>();
   if (!type || !type.hasStaticShape())
     return nullptr;

mlir/test/Dialect/Shape/canonicalize.mlir

 func @shape_eq_do_not_fold(%a : !shape.shape) -> i1 {
   // CHECK: %[[B:.*]] = shape.const_shape [4, 5, 6]
   // CHECK: %[[RESULT:.*]] = shape.shape_eq %[[A]], %[[B]] : !shape.shape, !shape.shape
   // CHECK: return %[[RESULT]] : i1
   %b = shape.const_shape [4, 5, 6] : !shape.shape
   %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
   return %result : i1
 }
+
+// -----
+
+// Fold `mul` for constant sizes.
+// CHECK-LABEL: @fold_mul_size
+func @fold_mul_size() -> !shape.size {
+  // CHECK: %[[RESULT:.*]] = shape.const_size 6
+  // CHECK: return %[[RESULT]] : !shape.size
+  %c2 = shape.const_size 2
+  %c3 = shape.const_size 3
+  %result = shape.mul %c2, %c3 : !shape.size, !shape.size -> !shape.size
+  return %result : !shape.size
+}
+
+// -----
+
+// Fold `mul` for constant indices.
+// CHECK-LABEL: @fold_mul_index
+func @fold_mul_index() -> index {
+  // CHECK: %[[RESULT:.*]] = constant 6 : index
+  // CHECK: return %[[RESULT]] : index
+  %c2 = constant 2 : index
+  %c3 = constant 3 : index
+  %result = shape.mul %c2, %c3 : index, index -> index
+  return %result : index
+}
+
+// -----
+
+// Fold `mul` for mixed constants.
+// CHECK-LABEL: @fold_mul_mixed
+func @fold_mul_mixed() -> !shape.size {
+  // CHECK: %[[RESULT:.*]] = shape.const_size 6
+  // CHECK: return %[[RESULT]] : !shape.size
+  %c2 = shape.const_size 2
+  %c3 = constant 3 : index
+  %result = shape.mul %c2, %c3 : !shape.size, index -> !shape.size
+  return %result : !shape.size
+}
+