use form traits in PolyMulToNTT where possible

PiperOrigin-RevId: 895546447

Author: j2kun

lib/Dialect/Polynomial/IR/PolynomialOps.td

@@ -437,7 +437,7 @@ def Polynomial_YieldOp : Polynomial_Op<"yield", [Terminator, HasParent<"ApplyCoe
 }
 
 def Polynomial_ApplyCoefficientwiseOp : Polynomial_Op<"apply_coefficientwise", [
-    Pure, SingleBlock]> {
+    Pure, SingleBlock, FixedFormCoeff]> {
   let summary = "Apply a region to each coefficient of a polynomial.";
   let description = [{
     `polynomial.apply_coefficientwise` takes a polynomial and applies a series

lib/Dialect/Polynomial/Transforms/PolyMulToNTT.cpp

@@ -2,6 +2,7 @@
 
 #include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
 #include "lib/Dialect/Polynomial/IR/PolynomialOps.h"
+#include "lib/Dialect/Polynomial/IR/PolynomialTraits.h"
 #include "lib/Dialect/Polynomial/IR/PolynomialTypes.h"
 #include "lib/Dialect/Polynomial/Transforms/NTTSolver.h"
 #include "llvm/include/llvm/ADT/DenseMap.h"              // from @llvm-project
@@ -57,21 +58,30 @@ enum class OpFormClass {
 };
 
 OpFormClass opFormClass(Operation* op) {
+  // This special case must come first because it overrides the more general
+  // SameOperandsAndResultForm trait for MulOp.
+  if (op->hasTrait<FixedFormEval>() || isa<MulOp>(op)) {
+    return OpFormClass::EVAL;
+  }
+
   if (isa<func::ReturnOp>(op)) {
     return OpFormClass::RETURN;
-  } else if (isa<ToTensorOp, LeadingTermOp, EvalOp, ConvertBasisOp,
-                 MonicMonomialMulOp, FromTensorOp, ApplyCoefficientwiseOp>(
-                 op)) {
+  }
+
+  if (op->hasTrait<FixedFormCoeff>()) {
     return OpFormClass::COEFF;
-  } else if (isa<MulOp>(op)) {
-    return OpFormClass::EVAL;
-  } else if (isa<AddOp, SubOp, MulScalarOp, ModSwitchOp, ExtractSliceOp,
-                 tensor::ExtractSliceOp, tensor::ExtractOp,
-                 tensor::FromElementsOp>(op)) {
+  }
+
+  if (op->hasTrait<SameOperandsAndResultForm>() ||
+      isa<ToTensorOp, FromTensorOp, tensor::ExtractSliceOp, tensor::ExtractOp,
+          tensor::FromElementsOp>(op)) {
     return OpFormClass::EITHER;
-  } else if (isa<MonomialOp, ConstantOp>(op)) {
+  }
+
+  if (isa<MonomialOp, ConstantOp>(op)) {
     return OpFormClass::CONST;
   }
+
   return OpFormClass::UNKNOWN;
 }