Add a folder-based EDSC intrinsics constructor (NFC)

PiperOrigin-RevId: 255908660

mlir/include/mlir/EDSC/Builders.h

@@ -26,6 +26,7 @@
 #include "mlir/AffineOps/AffineOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/StandardOps/Ops.h"
+#include "mlir/Transforms/FoldUtils.h"
 #include "mlir/VectorOps/VectorOps.h"
 
 namespace mlir {
@@ -315,6 +316,11 @@ public:
   template <typename Op, typename... Args>
   static ValueHandle create(Args... args);
 
+  /// Generic mlir::Op create. This is the key to being extensible to the whole
+  /// of MLIR without duplicating the type system or the op definitions.
+  template <typename Op, typename... Args>
+  static ValueHandle create(OperationFolder &folder, Args... args);
+
   /// Special case to build composed AffineApply operations.
   // TODO: createOrFold when available and move inside of the `create` method.
   static ValueHandle createComposedAffineApply(AffineMap map,
@@ -460,6 +466,12 @@ ValueHandle ValueHandle::create(Args... args) {
   llvm_unreachable("unsupported operation, use an OperationHandle instead");
 }
 
+template <typename Op, typename... Args>
+ValueHandle ValueHandle::create(OperationFolder &folder, Args... args) {
+  return ValueHandle(folder.create<Op>(ScopedContext::getBuilder(),
+                                       ScopedContext::getLocation(), args...));
+}
+
 namespace op {
 
 ValueHandle operator+(ValueHandle lhs, ValueHandle rhs);

mlir/include/mlir/EDSC/Intrinsics.h

@@ -118,6 +118,7 @@ inline detail::ValueHandleArray unpack(ArrayRef<ValueHandle> values) {
 /// Without subclassing, implicit conversion to Value* would fail when composing
 /// in patterns such as: `select(a, b, select(c, d, e))`.
 template <typename Op> struct ValueBuilder : public ValueHandle {
+  // Builder-based
   template <typename... Args>
   ValueBuilder(Args... args)
       : ValueHandle(ValueHandle::create<Op>(detail::unpack(args)...)) {}
@@ -136,6 +137,30 @@ template <typename Op> struct ValueBuilder : public ValueHandle {
       : ValueHandle(ValueHandle::create<Op>(
             detail::unpack(t1), detail::unpack(t2), detail::unpack(vs),
             detail::unpack(args)...)) {}
+
+  /// Folder-based
+  template <typename... Args>
+  ValueBuilder(OperationFolder &folder, Args... args)
+      : ValueHandle(ValueHandle::create<Op>(folder, detail::unpack(args)...)) {}
+  ValueBuilder(OperationFolder &folder, ArrayRef<ValueHandle> vs)
+      : ValueBuilder(ValueBuilder::create<Op>(folder, detail::unpack(vs))) {}
+  template <typename... Args>
+  ValueBuilder(OperationFolder &folder, ArrayRef<ValueHandle> vs, Args... args)
+      : ValueHandle(ValueHandle::create<Op>(folder, detail::unpack(vs),
+                                            detail::unpack(args)...)) {}
+  template <typename T, typename... Args>
+  ValueBuilder(OperationFolder &folder, T t, ArrayRef<ValueHandle> vs,
+               Args... args)
+      : ValueHandle(ValueHandle::create<Op>(folder, detail::unpack(t),
+                                            detail::unpack(vs),
+                                            detail::unpack(args)...)) {}
+  template <typename T1, typename T2, typename... Args>
+  ValueBuilder(OperationFolder &folder, T1 t1, T2 t2, ArrayRef<ValueHandle> vs,
+               Args... args)
+      : ValueHandle(ValueHandle::create<Op>(
+            folder, detail::unpack(t1), detail::unpack(t2), detail::unpack(vs),
+            detail::unpack(args)...)) {}
+
   ValueBuilder() : ValueHandle(ValueHandle::create<Op>()) {}
 };
 
@@ -162,15 +187,18 @@ template <typename Op> struct OperationBuilder : public OperationHandle {
 };
 
 using alloc = ValueBuilder<AllocOp>;
+using affine_apply = ValueBuilder<AffineApplyOp>;
 using constant_float = ValueBuilder<ConstantFloatOp>;
 using constant_index = ValueBuilder<ConstantIndexOp>;
 using constant_int = ValueBuilder<ConstantIntOp>;
 using dealloc = OperationBuilder<DeallocOp>;
 using dim = ValueBuilder<DimOp>;
 using load = ValueBuilder<LoadOp>;
+using muli = ValueBuilder<MulIOp>;
 using ret = OperationBuilder<ReturnOp>;
 using select = ValueBuilder<SelectOp>;
 using store = OperationBuilder<StoreOp>;
+using subi = ValueBuilder<SubIOp>;
 using vector_type_cast = ValueBuilder<VectorTypeCastOp>;
 
 /// Branches into the mlir::Block* captured by BlockHandle `b` with `operands`.

mlir/lib/EDSC/CMakeLists.txt

@@ -12,5 +12,6 @@ target_link_libraries(MLIREDSC
   PUBLIC
     MLIRAffineOps
     MLIRStandardOps
+    MLIRTransformUtils
     MLIRVectorOps
     )

mlir/lib/Linalg/IR/LinalgOps.cpp

@@ -869,8 +869,7 @@ foldedAffineApplies(OpBuilder &b, Location loc, AffineMap map,
     auto exprMap = AffineMap::get(dims, 0, e);
     SmallVector<Value *, 4> operands(vals.begin(), vals.end());
     canonicalizeMapAndOperands(&exprMap, &operands);
-    res.push_back(
-        ValueHandle(folder.create<AffineApplyOp>(b, loc, exprMap, operands)));
+    res.push_back(affine_apply(folder, exprMap, operands));
   }
   return res;
 }

mlir/lib/Transforms/CMakeLists.txt

@@ -1,6 +1,7 @@
+add_subdirectory(Utils)
+
 add_llvm_library(MLIRTransforms
   Canonicalizer.cpp
-  CMakeLists.txt
   CSE.cpp
   DialectConversion.cpp
   DmaGeneration.cpp
@@ -16,22 +17,18 @@ add_llvm_library(MLIRTransforms
   PipelineDataTransfer.cpp
   SimplifyAffineStructures.cpp
   StripDebugInfo.cpp
-  Utils/FoldUtils.cpp
-  Utils/GreedyPatternRewriteDriver.cpp
-  Utils/LoopFusionUtils.cpp
-  Utils/LoopUtils.cpp
-  Utils/RegionUtils.cpp
-  Utils/Utils.cpp
   Vectorize.cpp
   ViewFunctionGraph.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
   )
+
 add_dependencies(MLIRTransforms MLIRStandardOpsIncGen)
 target_link_libraries(MLIRTransforms
   MLIRAffineOps
   MLIRAnalysis
   MLIRPass
+  MLIRTransformUtils
   MLIRVectorOps
   )

mlir/lib/Transforms/Utils/CMakeLists.txt

@@ -0,0 +1,18 @@
+add_llvm_library(MLIRTransformUtils
+  FoldUtils.cpp
+  GreedyPatternRewriteDriver.cpp
+  LoopFusionUtils.cpp
+  LoopUtils.cpp
+  RegionUtils.cpp
+  Utils.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
+  )
+
+add_dependencies(MLIRTransformUtils MLIRStandardOpsIncGen)
+target_link_libraries(MLIRTransformUtils
+  MLIRAffineOps
+  MLIRAnalysis
+  MLIRPass
+  )