[mlir][spirv] Convert linalg.generic for reduction to SPIR-V ops

This commit adds a pattern to lower linalg.generic for reduction
to spv.GroupNonUniform* ops. Right now this only supports integer
reduction on 1-D input memref. Shader entry point ABI is queried
to make sure that the input memref's shape matches the local
workgroup's invocation configuration. This makes sure that the
workload fits in one local workgroup so that we can leverage
SPIR-V group non-uniform operations.

linglg.generic is a structured op that preserves the right level
of information. It is easier to recognize reduction at this level
than performing analysis on loops.

This commit also exposes `getElementPtr` in SPIRVLowering.h given
that it's a generally useful utility function.

Differential Revision: https://reviews.llvm.org/D73437

Author: antiagainst

mlir/include/mlir/Conversion/LinalgToSPIRV/LinalgToSPIRV.h

@@ -0,0 +1,29 @@
+//===- LinalgToSPIRV.h - Linalg to SPIR-V dialect conversion ----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file provides patterns for Linalg to SPIR-V dialect conversion.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_LINALGTOSPIRV_LINALGTOSPIRV_H
+#define MLIR_CONVERSION_LINALGTOSPIRV_LINALGTOSPIRV_H
+
+namespace mlir {
+class MLIRContext;
+class OwningRewritePatternList;
+class SPIRVTypeConverter;
+
+/// Appends to a pattern list additional patterns for translating Linalg ops to
+/// SPIR-V ops.
+void populateLinalgToSPIRVPatterns(MLIRContext *context,
+                                   SPIRVTypeConverter &typeConverter,
+                                   OwningRewritePatternList &patterns);
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_LINALGTOSPIRV_LINALGTOSPIRV_H

mlir/include/mlir/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.h

@@ -0,0 +1,25 @@
+//===- LinalgToSPIRVPass.h -  Linalg to SPIR-V conversion pass --*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file provides a pass for Linalg to SPIR-V dialect conversion.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_STANDARDTOSPIRV_LINALGTOSPIRVPASS_H
+#define MLIR_CONVERSION_STANDARDTOSPIRV_LINALGTOSPIRVPASS_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+/// Creates and returns a pass to convert Linalg ops to SPIR-V ops.
+std::unique_ptr<OpPassBase<ModuleOp>> createLinalgToSPIRVPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_STANDARDTOSPIRV_LINALGTOSPIRVPASS_H

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

@@ -54,6 +54,11 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
       "Query the number of loops within the current operation.",
       "unsigned", "getNumLoops">,
 
+    InterfaceMethod<
+      [{Returns true if the current operation has only one loop and it's a
+        reduction loop}],
+      "unsigned", "hasSingleReductionLoop">,
+
     //========================================================================//
     // Input arguments handling.
     //========================================================================//

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

@@ -93,6 +93,12 @@ class StructuredOpTraits
         cast<ConcreteType>(this->getOperation()).iterator_types());
   }
 
+  bool hasSingleReductionLoop() {
+    auto iterators = cast<ConcreteType>(this->getOperation()).iterator_types();
+    return iterators.size() == 1 &&
+           getNumIterators(getReductionIteratorTypeName(), iterators);
+  }
+
   //==========================================================================//
   // Input arguments handling.
   //==========================================================================//

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

@@ -28,6 +28,26 @@ struct FusionInfo {
   LinalgOp fusedProducer;
 };
 
+/// A struct containing common matchers over linalg op's region.
+struct RegionMatcher {
+  enum class BinaryOpKind {
+    IAdd,
+  };
+
+  /// Matches the given linalg op if its body is performing binary operation on
+  /// int or float scalar values and returns the binary op kind.
+  ///
+  /// The linalg op's region is expected to be
+  /// ```
+  /// {
+  ///   ^bb(%a: <scalar-type>, %b: <scalar-type>):
+  ///     %0 = <binary-op> %a, %b: <scalar-type>
+  ///     linalg.yield %0: <scalar-type>
+  /// }
+  /// ```
+  static Optional<BinaryOpKind> matchAsScalarBinaryOp(GenericOp op);
+};
+
 /// Checks whether the specific `producer` is the last write to exactly the
 /// whole `consumedView`. This checks structural dominance, that the dependence
 /// is a RAW without any interleaved write to any piece of `consumedView`.

mlir/include/mlir/Dialect/SPIRV/SPIRVAtomicOps.td

@@ -25,6 +25,7 @@ class SPV_AtomicUpdateOp<string mnemonic, list<OpTrait> traits = []> :
     SPV_ScopeAttr:$memory_scope,
     SPV_MemorySemanticsAttr:$semantics
   );
+
   let results = (outs
     SPV_Integer:$result
   );
@@ -42,9 +43,19 @@ class SPV_AtomicUpdateWithValueOp<string mnemonic, list<OpTrait> traits = []> :
     SPV_MemorySemanticsAttr:$semantics,
     SPV_Integer:$value
   );
+
   let results = (outs
     SPV_Integer:$result
   );
+
+  let builders = [
+    OpBuilder<
+      [{Builder *builder, OperationState &state, Value pointer,
+        ::mlir::spirv::Scope scope, ::mlir::spirv::MemorySemantics memory,
+        Value value}],
+      [{build(builder, state, value.getType(), pointer, scope, memory, value);}]
+    >
+  ];
 }
 
 // -----

mlir/include/mlir/Dialect/SPIRV/SPIRVControlFlowOps.td

@@ -453,14 +453,22 @@ def SPV_SelectionOp : SPV_Op<"selection", [InFunctionScope]> {
   let regions = (region AnyRegion:$body);
 
   let extraClassDeclaration = [{
-    // Returns the selection header block.
+    /// Returns the selection header block.
     Block *getHeaderBlock();
 
-    // Returns the selection merge block.
+    /// Returns the selection merge block.
     Block *getMergeBlock();
 
-    // Adds a selection merge block containing one spv._merge op.
+    /// Adds a selection merge block containing one spv._merge op.
     void addMergeBlock();
+
+    /// Creates a spv.selection op for `if (<condition>) then { <thenBody> }`
+    /// with `builder`. `builder`'s insertion point will remain at after the
+    /// newly inserted spv.selection op afterwards.
+    static SelectionOp createIfThen(
+        Location loc, Value condition,
+        function_ref<void(OpBuilder *builder)> thenBody,
+        OpBuilder *builder);
   }];
 
   let hasOpcode = 0;

mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h

@@ -58,6 +58,8 @@ void populateBuiltinFuncToSPIRVPatterns(MLIRContext *context,
                                         OwningRewritePatternList &patterns);
 
 namespace spirv {
+class AccessChainOp;
+
 class SPIRVConversionTarget : public ConversionTarget {
 public:
   /// Creates a SPIR-V conversion target for the given target environment.
@@ -90,6 +92,16 @@ class SPIRVConversionTarget : public ConversionTarget {
 Value getBuiltinVariableValue(Operation *op, BuiltIn builtin,
                               OpBuilder &builder);
 
+/// Performs the index computation to get to the element at `indices` of the
+/// memory pointed to by `basePtr`, using the layout map of `baseType`.
+
+// TODO(ravishankarm) : This method assumes that the `baseType` is a MemRefType
+// with AffineMap that has static strides. Extend to handle dynamic strides.
+spirv::AccessChainOp getElementPtr(SPIRVTypeConverter &typeConverter,
+                                   MemRefType baseType, Value basePtr,
+                                   ArrayRef<Value> indices, Location loc,
+                                   OpBuilder &builder);
+
 /// Sets the InterfaceVarABIAttr and EntryPointABIAttr for a function and its
 /// arguments.
 LogicalResult setABIAttrs(FuncOp funcOp, EntryPointABIAttr entryPointInfo,

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h

@@ -54,6 +54,12 @@ TargetEnvAttr getDefaultTargetEnv(MLIRContext *context);
 /// target environment (SPIR-V 1.0 with Shader capability and no extra
 /// extensions) if not provided.
 TargetEnvAttr lookupTargetEnvOrDefault(Operation *op);
+
+/// Queries the local workgroup size from entry point ABI on the nearest
+/// function-like op containing the given `op`. Returns null attribute if not
+/// found.
+DenseIntElementsAttr lookupLocalWorkGroupSize(Operation *op);
+
 } // namespace spirv
 } // namespace mlir
 

mlir/lib/Conversion/CMakeLists.txt

@@ -4,6 +4,7 @@ add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(LinalgToLLVM)
+add_subdirectory(LinalgToSPIRV)
 add_subdirectory(LoopsToGPU)
 add_subdirectory(LoopToStandard)
 add_subdirectory(StandardToLLVM)