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

[antiagainst]

Summary

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.

[merge_guards_bot]

Unit tests: pass. 62199 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

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

4

Should be LLVM Project now.

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

28

Why the inline keyword?

61

Don't use size, it is O(N). Use 'has_single_element' instead.

71

Same here, size() is O(N): !has_single_element(block.without_terminator())

104

This function is huge, is there any way to split it up a bit?

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.cpp

43

Drop the trivial brace.

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

70

nit: Drop the newline here.

[mravishankar]

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

37

Nit : s/interators/iterators/

74

I feel like this is not required. If the linalg.generic operation says this is a reduction, we should not be checking the region to verify that it is. linalg as a contract is guaranteeing this is a reduction.

183

This is hard-wiring to IAdd. We probably need to structure this differently. We need to have a pattern to lower the linalg.generic with reduction iterator into the kernel generated here, and then lower the operations within the region separately.

[antiagainst]

Address comments

[antiagainst]

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

28

Added initially because this function was trivial enough. I guess right now it's better to let the compiler decide.

61

Oh, good to know! Thanks!

71

Nice, thanks!

74

I'm not sure I get your idea correctly; but this is checking we are doing the expected kind of reduction (BinaryOp).

104

Good point. Split out two helper functions.

183

It was intentional to only support IAdd. I've re-structured this a bit so it's extensible to other binary op kinds.

[antiagainst]

Revert unrelated commits

[antiagainst]

Clean up unrelated commits again

[antiagainst]

Messed up the revision history with Arc... Please show the diff between "Diff 1" and "Diff 4" to check the modifications.

[merge_guards_bot]

Unit tests: pass. 62199 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: pass. 62199 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

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

470

Drop the llvm::.

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

33

nit: You can use getOperandTypes() && getResultTypes() respectively.

39

Can you use size() directly on the ArrayAttr?

67

nit: Missing end namespace comment.

81

You should be able to remove the .empty() call as well as the .getBlocks()

[nicolasvasilache]

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

28

Please use LinalgOp::hasBufferSemantics instead of rolling your own

38

Can we create a singleReductionIterator() somewhere in the existing utils and directly equality compare the ArrayAttr iterators against that?
I anticipate this type of things is soon going to become more and more useful, including from Tablegen.

79

Same thing here, we should move these to utils on the Linalg side and make reusable as things are soon going to blow up out of control otherwise.

126

Please split this into a precondition and an apply that must succeed so we can use the precondition as a Constraint and the apply as a simple Rewrite.
I have found this split to be crucial to write custom fused passes with DRR.
I claim this is generally the way we should be writing transformations when we can, which is the case for Linalg ops.

157

I have similar comments re utils here but it will take a bit longer to put things in a reasonable form so please just add a TODO that this should use generic Linalg utils when available.

261

Thanks for doing this as a pattern, it will compose nicely thanks to all of @rriddle's great work!

[mravishankar]

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

74

That is what I am not sure we need to do. You have already checked the generic op is a 1D loop with reduction iterator type. So the body of the generic op is assumed to be a reduction right, i.e. a "binaryOp". Here you are checking whether its an operation which takes two operands, but not necessarily a binary operation that can be used for reduction. In some sense you dont need to check that since the generic op description already told you to assume that the region of the op represents a binary op that can be used for reduction (this is based on my understanding of "reduction" loops in linalg, but I need to double check).

183

The way I see this structured is

1. You check the "structure" of the linalg.generic op to see if it is a 1D reduction. You assume that the body of the generic op represents a binary operation that can be used for the reduction.
2. You write a separate pattern that converts the operations of the body itself into the spirv::GroupNonUniform*Op.

The dialect conversion will first convert the generic op, and then it will attempt to convert the body of the op. The way this is strucutred, it can only handle straight-forward binary operations. THere could be more complicated operations in the region of a generic op that implements the reduction, which would be hard to incorporate in this approach.

With your updated changes, it is easy to extend to other ops, but I think a more general solution is needed where we are not constrained in handling just simple reductions. It might need some more careful design though, which I dont have a full picture of right now. So for now this OK, but please leave a TODO to say something like "handle more general reductions".

[mravishankar]

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

74

Ok, looked into this a little bit more. This is indeed required as of now. Ideally we dont need to do this, but that is orthogonal. No issues here.

183

Thought about this a little bit more. It will be a more sane strategy to not handle cases where the region is "inlined" into the generic op, but rather the region is specified as a function using the fun attribute of the linalg.generic op. Then we can split the lowering of the generic op and the lowering of the function that does the reduction separately.
This too not needed for this CL. Should revisit this

[antiagainst]

Address comments

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

28

Oh, didn't know we have that. Thanks!

81

Nice, thanks!

126

I feel I don't have all the backgrounds here so I'd appreciate some explanation as for why it's necessary to separate the match and rewrite. I think that's the original design of patterns (that we have match and rewrite separately) and then later we found it's too cumbersome to carry over a large matched state between them and then came up with matchAndRewrite? IIUC, every pattern is applied as a whole. I can understand that the match/rewrite side can be potentially reusable across patterns as components if they are exposed as helper functions, but that's more from a library organization and code reuse's perspective; I feel you've more reasons than that so I'd like to learn about them. :)

Another question here is that I've matches that is generally useful to other patterns (checking whether this is a linalg doing reduction) and matches that is only useful to SPIR-V here (checking whether the workload can be fit in one local workgroup). How do you handle such cases?

For now I separated the check that whether this is a linalg doing reduction as a helper function. I haven't put it in some linalg file yet because currently it's too rigid and overfitting to a specific case. I plan to extend it and then later maybe we can move it to some linalg file.

157

Right. I guess here we need some thinking about how to better design the API to make it generic and usable across different cases. I'm not so versed on that at the current moment. So putting your name in the TODO for now. :)

183

Yeah, right now linalg.generic is an integrated entity; I'm not sure how we can easily convert part of it (the "structure") and then the inlined scalar region. I guess it's fine for now to overfit a bit until later we see concrete cases that we want to support, which will give us a more tangible problem to solve.

[merge_guards_bot]

Unit tests: pass. 62199 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[nicolasvasilache]

Thanks Lei, this looks great, glad to see you pushing on this front!

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

126

Let's sync up over video for this, this is a longer discussion and touches phase ordering issues, local patterns etc.
If you want some of the deeper details please look at the rationale doc: https://mlir.llvm.org/docs/Dialects/Linalg/ (which will soon move to https://mlir.llvm.org/docs/Dialects/LinalgRationale/).

In the meantime, this is not a blocker but it would be great to sync up in particular on the implementation on those ideas today in DRR.

157

SG, thanks!

[antiagainst]

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

126

Awesome, thanks! I've opened https://mlir.llvm.org/docs/Dialects/Linalg/ in a tab but haven't read it through yet. ;-P Will do now.

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

+//===- 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

+//===- 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

let methods = [
     InterfaceMethod<
       "Query the number of window loops within the current operation.",
       "unsigned", "getNumWindowLoops"
     >,
     InterfaceMethod<
       "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.
     //========================================================================//
     InterfaceMethod<
       "Query the number of inputs from the current operation.",
       "unsigned", "getNumInputs"
     >,
     InterfaceMethod<"Query the input view at the given index.",

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

return getNumIterators(
         getWindowIteratorTypeName(),
         cast<ConcreteType>(this->getOperation()).iterator_types());
   }
   unsigned getNumLoops() {
     return getNumIterators(
         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.
   //==========================================================================//
   // The `i^th` input argument is always the `i^th` operand regardless of
   // whether we have tensors or buffers.
   //
   /// Return the `i`-th input value.
   Value getInput(unsigned i) {

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

 namespace linalg {
 class LinalgDependenceGraph;
 
 struct FusionInfo {
   LinalgOp originalProducer;
   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`.
 bool isProducerLastWriteOfView(const LinalgDependenceGraph &graph,
                                LinalgOp consumer, Value consumedView,
                                LinalgOp producer);
 
 /// Checks whether fusing the specific `producer` of the `consumedView` is

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

class SPV_AtomicUpdateOp<string mnemonic, list<OpTrait> traits = []> :
   let printer = [{ return ::printAtomicUpdateOp(getOperation(), p); }];
   let verifier = [{ return ::verifyAtomicUpdateOp(getOperation()); }];
 
   let arguments = (ins
     SPV_AnyPtr:$pointer,
     SPV_ScopeAttr:$memory_scope,
     SPV_MemorySemanticsAttr:$semantics
   );
+
   let results = (outs
     SPV_Integer:$result
   );
 }
 
 class SPV_AtomicUpdateWithValueOp<string mnemonic, list<OpTrait> traits = []> :
   SPV_Op<mnemonic, traits> {
   let parser = [{ return ::parseAtomicUpdateOp(parser, result, true); }];
   let printer = [{ return ::printAtomicUpdateOp(getOperation(), p); }];
   let verifier = [{ return ::verifyAtomicUpdateOp(getOperation()); }];
 
   let arguments = (ins
     SPV_AnyPtr:$pointer,
     SPV_ScopeAttr:$memory_scope,
     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);}]
+    >
+  ];
 }
 
 // -----
 
 def SPV_AtomicAndOp : SPV_AtomicUpdateWithValueOp<"AtomicAnd", []> {
   let summary = [{
     Perform the following steps atomically with respect to any other atomic
     accesses within Scope to the same location:

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

let arguments = (ins
     SPV_SelectionControlAttr:$selection_control
   );
 
   let results = (outs);
 
   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;
 
   let autogenSerialization = 0;
 
   let hasCanonicalizer = 1;
 }
 
 #endif // SPIRV_CONTROLFLOW_OPS

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

 /// `func` op to the SPIR-V dialect. These patterns do not handle shader
 /// interface/ABI; they convert function parameters to be of SPIR-V allowed
 /// types.
 void populateBuiltinFuncToSPIRVPatterns(MLIRContext *context,
                                         SPIRVTypeConverter &typeConverter,
                                         OwningRewritePatternList &patterns);
 
 namespace spirv {
+class AccessChainOp;
+
 class SPIRVConversionTarget : public ConversionTarget {
 public:
   /// Creates a SPIR-V conversion target for the given target environment.
   static std::unique_ptr<SPIRVConversionTarget> get(TargetEnvAttr targetEnv,
                                                     MLIRContext *context);
 
 private:
   SPIRVConversionTarget(TargetEnvAttr targetEnv, MLIRContext *context);

 
 /// Returns the value for the given `builtin` variable. This function gets or
 /// inserts the global variable associated for the builtin within the nearest
 /// enclosing op that has a symbol table. Returns null Value if such an
 /// enclosing op cannot be found.
 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,
                           ArrayRef<InterfaceVarABIAttr> argABIInfo);
 } // namespace spirv
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_SPIRVLOWERING_H

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

 /// Returns the default target environment: SPIR-V 1.0 with Shader capability
 /// and no extra extensions.
 TargetEnvAttr getDefaultTargetEnv(MLIRContext *context);
 
 /// Queries the target environment from the given `op` or returns the default
 /// 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
 
 #endif // MLIR_DIALECT_SPIRV_TARGETANDABI_H

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(GPUToCUDA)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(LinalgToLLVM)
+add_subdirectory(LinalgToSPIRV)
 add_subdirectory(LoopsToGPU)
 add_subdirectory(LoopToStandard)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToLoops)

mlir/lib/Conversion/LinalgToSPIRV/CMakeLists.txt

+add_llvm_library(MLIRLinalgToSPIRVTransforms
+  LinalgToSPIRV.cpp
+  LinalgToSPIRVPass.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SPIRV
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
+  )
+
+target_link_libraries(MLIRLinalgToSPIRVTransforms
+  MLIRIR
+  MLIRLinalgOps
+  MLIRPass
+  MLIRSPIRV
+  MLIRSupport
+  )

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

+//===- LinalgToSPIRV.cpp - Linalg to SPIR-V dialect conversion ------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRV.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/SPIRV/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/SPIRVLowering.h"
+#include "mlir/Dialect/SPIRV/SPIRVOps.h"
+#include "mlir/Dialect/StandardOps/Ops.h"
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
+#include "mlir/IR/AffineExpr.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Utilities
+//===----------------------------------------------------------------------===//
+
+/// Returns a `Value` containing the `dim`-th dimension's size of SPIR-V
+/// location invocation ID. This function will create necessary operations with
+/// `builder` at the proper region containing `op`.
+static Value getLocalInvocationDimSize(Operation *op, int dim, Location loc,
+                                       OpBuilder *builder) {
+  assert(dim >= 0 && dim < 3 && "local invocation only has three dimensions");
+  Value invocation = spirv::getBuiltinVariableValue(
+      op, spirv::BuiltIn::LocalInvocationId, *builder);
+  Type xType = invocation.getType().cast<ShapedType>().getElementType();
+  return builder->create<spirv::CompositeExtractOp>(
+      loc, xType, invocation, builder->getI32ArrayAttr({dim}));
+}
+
+
+//===----------------------------------------------------------------------===//
+// Reduction (single workgroup)
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// A pattern to convert a linalg.generic op to SPIR-V ops under the condition
+/// that the linalg.generic op is performing reduction with a workload size that
+/// can fit in one workgroup.
+class SingleWorkgroupReduction final
+    : public SPIRVOpLowering<linalg::GenericOp> {
+public:
+  using SPIRVOpLowering<linalg::GenericOp>::SPIRVOpLowering;
+
+  /// Matches the given linalg.generic op as performing reduction and returns
+  /// the binary op kind if successful.
+  static Optional<linalg::RegionMatcher::BinaryOpKind>
+  matchAsPerformingReduction(linalg::GenericOp genericOp);
+
+  PatternMatchResult
+  matchAndRewrite(linalg::GenericOp genericOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
+} // namespace
+
+Optional<linalg::RegionMatcher::BinaryOpKind>
+SingleWorkgroupReduction::matchAsPerformingReduction(
+    linalg::GenericOp genericOp) {
+  Operation *op = genericOp.getOperation();
+
+  // Make sure the linalg.generic is working on memrefs.
+  if (!genericOp.hasBufferSemantics())
+    return llvm::None;
+
+  // Make sure this is reudction with one input and one output.
+  if (genericOp.args_in().getZExtValue() != 1 ||
+      genericOp.args_out().getZExtValue() != 1)
+    return llvm::None;
+
+  auto originalInputType = op->getOperand(0).getType().cast<MemRefType>();
+  auto originalOutputType = op->getOperand(1).getType().cast<MemRefType>();
+
+  // Make sure the original input has one dimension.
+  if (!originalInputType.hasStaticShape() || originalInputType.getRank() != 1)
+    return llvm::None;
+  // Make sure the original output has one element.
+  if (!originalOutputType.hasStaticShape() ||
+      originalOutputType.getNumElements() != 1)
+    return llvm::None;
+
+  if (!genericOp.hasSingleReductionLoop())
+    return llvm::None;
+
+  if (genericOp.indexing_maps().getValue().size() != 2)
+    return llvm::None;
+
+  // TODO(nicolasvasilache): create utility functions for these checks in Linalg
+  // and use them.
+  auto inputMap = genericOp.indexing_maps().getValue()[0].cast<AffineMapAttr>();
+  auto outputMap =
+      genericOp.indexing_maps().getValue()[1].cast<AffineMapAttr>();
+  // The indexing map for the input should be `(i) -> (i)`.
+  if (inputMap.getValue() !=
+      AffineMap::get(1, 0, {getAffineDimExpr(0, op->getContext())}))
+    return llvm::None;
+  // The indexing map for the input should be `(i) -> (0)`.
+  if (outputMap.getValue() !=
+      AffineMap::get(1, 0, {getAffineConstantExpr(0, op->getContext())}))
+    return llvm::None;
+
+  return linalg::RegionMatcher::matchAsScalarBinaryOp(genericOp);
+}
+
+PatternMatchResult SingleWorkgroupReduction::matchAndRewrite(
+    linalg::GenericOp genericOp, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+  Operation *op = genericOp.getOperation();
+  auto originalInputType = op->getOperand(0).getType().cast<MemRefType>();
+  auto originalOutputType = op->getOperand(1).getType().cast<MemRefType>();
+
+  auto binaryOpKind = matchAsPerformingReduction(genericOp);
+  if (!binaryOpKind)
+    return matchFailure();
+
+  // Query the shader interface for local workgroup size to make sure the
+  // invocation configuration fits with the input memref's shape.
+  DenseIntElementsAttr localSize = spirv::lookupLocalWorkGroupSize(genericOp);
+  if (!localSize)
+    return matchFailure();
+
+  if ((*localSize.begin()).getSExtValue() != originalInputType.getDimSize(0))
+    return matchFailure();
+  if (llvm::any_of(llvm::drop_begin(localSize.getIntValues(), 1),
+                   [](const APInt &size) { return !size.isOneValue(); }))
+    return matchFailure();
+
+  // TODO(antiagainst): Query the target environment to make sure the current
+  // workload fits in a local workgroup.
+
+  Value convertedInput = operands[0], convertedOutput = operands[1];
+  Location loc = genericOp.getLoc();
+
+  // Get the invocation ID.
+  Value x = getLocalInvocationDimSize(genericOp, /*dim=*/0, loc, &rewriter);
+
+  // TODO(antiagainst): Load to Workgroup storage class first.
+
+  // Get the input element accessed by this invocation.
+  Value inputElementPtr = spirv::getElementPtr(
+      typeConverter, originalInputType, convertedInput, {x}, loc, rewriter);
+  Value inputElement = rewriter.create<spirv::LoadOp>(loc, inputElementPtr);
+
+  // Perform the group reduction operation.
+  Value groupOperation;
+#define CREATE_GROUP_NON_UNIFORM_BIN_OP(opKind, spvOp)                         \
+  case linalg::RegionMatcher::BinaryOpKind::opKind: {                          \
+    groupOperation = rewriter.create<spirv::spvOp>(                            \
+        loc, originalInputType.getElementType(), spirv::Scope::Subgroup,       \
+        spirv::GroupOperation::Reduce, inputElement,                           \
+        /*cluster_size=*/ArrayRef<Value>());                                   \
+  } break
+  switch (*binaryOpKind) {
+    CREATE_GROUP_NON_UNIFORM_BIN_OP(IAdd, GroupNonUniformIAddOp);
+  }
+#undef CREATE_GROUP_NON_UNIFORM_BIN_OP
+
+  // Get the output element accessed by this reduction.
+  Value zero = spirv::ConstantOp::getZero(
+      typeConverter.getIndexType(rewriter.getContext()), loc, &rewriter);
+  SmallVector<Value, 1> zeroIndices(originalOutputType.getRank(), zero);
+  Value outputElementPtr =
+      spirv::getElementPtr(typeConverter, originalOutputType, convertedOutput,
+                           zeroIndices, loc, rewriter);
+
+  // Write out the final reduction result. This should be only conducted by one
+  // invocation. We use spv.GroupNonUniformElect to find the invocation with the
+  // lowest ID.
+  //
+  // ```
+  // if (spv.GroupNonUniformElect) { output = ... }
+  // ```
+
+  Value condition = rewriter.create<spirv::GroupNonUniformElectOp>(
+      loc, spirv::Scope::Subgroup);
+
+  auto createAtomicOp = [&](OpBuilder *builder) {
+#define CREATE_ATOMIC_BIN_OP(opKind, spvOp)                                    \
+  case linalg::RegionMatcher::BinaryOpKind::opKind: {                          \
+    builder->create<spirv::spvOp>(loc, outputElementPtr, spirv::Scope::Device, \
+                                  spirv::MemorySemantics::AcquireRelease,      \
+                                  groupOperation);                             \
+  } break
+    switch (*binaryOpKind) { CREATE_ATOMIC_BIN_OP(IAdd, AtomicIAddOp); }
+#undef CREATE_ATOMIC_BIN_OP
+  };
+
+  spirv::SelectionOp::createIfThen(loc, condition, createAtomicOp, &rewriter);
+
+  rewriter.eraseOp(genericOp);
+  return matchSuccess();
+}
+
+//===----------------------------------------------------------------------===//
+// Pattern population
+//===----------------------------------------------------------------------===//
+
+void mlir::populateLinalgToSPIRVPatterns(MLIRContext *context,
+                                         SPIRVTypeConverter &typeConverter,
+                                         OwningRewritePatternList &patterns) {
+  patterns.insert<SingleWorkgroupReduction>(context, typeConverter);
+}

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.cpp

+//===- LinalgToSPIRVPass.cpp - Linalg to SPIR-V conversion pass -----------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.h"
+#include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRV.h"
+#include "mlir/Dialect/SPIRV/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/SPIRVLowering.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+/// A pass converting MLIR Linalg ops into SPIR-V ops.
+class LinalgToSPIRVPass : public ModulePass<LinalgToSPIRVPass> {
+  void runOnModule() override;
+};
+} // namespace
+
+void LinalgToSPIRVPass::runOnModule() {
+  MLIRContext *context = &getContext();
+  ModuleOp module = getModule();
+
+  SPIRVTypeConverter typeConverter;
+  OwningRewritePatternList patterns;
+  populateLinalgToSPIRVPatterns(context, typeConverter, patterns);
+  populateBuiltinFuncToSPIRVPatterns(context, typeConverter, patterns);
+
+  auto targetEnv = spirv::lookupTargetEnvOrDefault(module);
+  std::unique_ptr<ConversionTarget> target =
+      spirv::SPIRVConversionTarget::get(targetEnv, context);
+
+  // Allow builtin ops.
+  target->addLegalOp<ModuleOp, ModuleTerminatorOp>();
+  target->addDynamicallyLegalOp<FuncOp>(
+      [&](FuncOp op) { return typeConverter.isSignatureLegal(op.getType()); });
+
+  if (failed(applyFullConversion(module, *target, patterns)))
+    return signalPassFailure();
+}
+
+std::unique_ptr<OpPassBase<ModuleOp>> mlir::createLinalgToSPIRVPass() {
+  return std::make_unique<LinalgToSPIRVPass>();
+}
+
+static PassRegistration<LinalgToSPIRVPass>
+    pass("convert-linalg-to-spirv", "Convert Linalg ops to SPIR-V ops");

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

public:
   PatternMatchResult
   matchAndRewrite(StoreOp storeOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// Utility functions for operation conversion
-//===----------------------------------------------------------------------===//
-
-/// Performs the index computation to get to the element pointed to by
-/// `indices` using the layout map of `baseType`.
-
-// TODO(ravishankarm) : This method assumes that the `origBaseType` is a
-// MemRefType with AffineMap that has static strides. Handle dynamic strides
-static spirv::AccessChainOp getElementPtr(OpBuilder &builder,
-                                          SPIRVTypeConverter &typeConverter,
-                                          Location loc, MemRefType origBaseType,
-                                          Value basePtr,
-                                          ArrayRef<Value> indices) {
-  // Get base and offset of the MemRefType and verify they are static.
-  int64_t offset;
-  SmallVector<int64_t, 4> strides;
-  if (failed(getStridesAndOffset(origBaseType, strides, offset)) ||
-      llvm::is_contained(strides, MemRefType::getDynamicStrideOrOffset())) {
-    return nullptr;
-  }
-
-  auto indexType = typeConverter.getIndexType(builder.getContext());
-
-  Value ptrLoc = nullptr;
-  assert(indices.size() == strides.size());
-  for (auto index : enumerate(indices)) {
-    Value strideVal = builder.create<spirv::ConstantOp>(
-        loc, indexType, IntegerAttr::get(indexType, strides[index.index()]));
-    Value update = builder.create<spirv::IMulOp>(loc, strideVal, index.value());
-    ptrLoc =
-        (ptrLoc ? builder.create<spirv::IAddOp>(loc, ptrLoc, update).getResult()
-                : update);
-  }
-  SmallVector<Value, 2> linearizedIndices;
-  // Add a '0' at the start to index into the struct.
-  linearizedIndices.push_back(builder.create<spirv::ConstantOp>(
-      loc, indexType, IntegerAttr::get(indexType, 0)));
-  linearizedIndices.push_back(ptrLoc);
-  return builder.create<spirv::AccessChainOp>(loc, basePtr, linearizedIndices);
-}
-
-//===----------------------------------------------------------------------===//
 // ConstantOp with composite type.
 //===----------------------------------------------------------------------===//
 
 PatternMatchResult ConstantCompositeOpConversion::matchAndRewrite(
     ConstantOp constCompositeOp, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   auto compositeType =
       constCompositeOp.getResult().getType().dyn_cast<RankedTensorType>();

 //===----------------------------------------------------------------------===//
 // LoadOp
 //===----------------------------------------------------------------------===//
 
 PatternMatchResult
 LoadOpConversion::matchAndRewrite(LoadOp loadOp, ArrayRef<Value> operands,
                                   ConversionPatternRewriter &rewriter) const {
   LoadOpOperandAdaptor loadOperands(operands);
-  auto loadPtr = getElementPtr(rewriter, typeConverter, loadOp.getLoc(),
-                               loadOp.memref().getType().cast<MemRefType>(),
-                               loadOperands.memref(), loadOperands.indices());
+  auto loadPtr = spirv::getElementPtr(
+      typeConverter, loadOp.memref().getType().cast<MemRefType>(),
+      loadOperands.memref(), loadOperands.indices(), loadOp.getLoc(), rewriter);
   rewriter.replaceOpWithNewOp<spirv::LoadOp>(loadOp, loadPtr);
   return matchSuccess();
 }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 

 //===----------------------------------------------------------------------===//
 // StoreOp
 //===----------------------------------------------------------------------===//
 
 PatternMatchResult
 StoreOpConversion::matchAndRewrite(StoreOp storeOp, ArrayRef<Value> operands,
                                    ConversionPatternRewriter &rewriter) const {
   StoreOpOperandAdaptor storeOperands(operands);
-  auto storePtr =
-      getElementPtr(rewriter, typeConverter, storeOp.getLoc(),
-                    storeOp.memref().getType().cast<MemRefType>(),
-                    storeOperands.memref(), storeOperands.indices());
+  auto storePtr = spirv::getElementPtr(
+      typeConverter, storeOp.memref().getType().cast<MemRefType>(),
+      storeOperands.memref(), storeOperands.indices(), storeOp.getLoc(),
+      rewriter);
   rewriter.replaceOpWithNewOp<spirv::StoreOp>(storeOp, storePtr,
                                               storeOperands.value());
   return matchSuccess();
 }
 
 namespace {
 /// Import the Standard Ops to SPIR-V Patterns.
 #include "StandardToSPIRV.cpp.inc"

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/EDSC/Helpers.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/STLExtras.h"
 #include "mlir/Transforms/FoldUtils.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 using namespace mlir::linalg;
 using namespace mlir::loop;
 
+Optional<RegionMatcher::BinaryOpKind>
+RegionMatcher::matchAsScalarBinaryOp(GenericOp op) {
+  auto &region = op.region();
+  if (!has_single_element(region))
+    return llvm::None;
+
+  Block &block = region.front();
+  if (block.getNumArguments() != 2 ||
+      !block.getArgument(0).getType().isIntOrFloat() ||
+      !block.getArgument(1).getType().isIntOrFloat())
+    return llvm::None;
+
+  auto &ops = block.getOperations();
+  if (!has_single_element(block.without_terminator()))
+    return llvm::None;
+
+  using mlir::matchers::m_Val;
+  auto a = m_Val(block.getArgument(0));
+  auto b = m_Val(block.getArgument(1));
+
+  auto addPattern = m_Op<linalg::YieldOp>(m_Op<AddIOp>(a, b));
+  if (addPattern.match(&ops.back()))
+    return BinaryOpKind::IAdd;
+
+  return llvm::None;
+}
+
 static Value emitOrFoldComposedAffineApply(OpBuilder &b, Location loc,
                                            AffineMap map,
                                            ArrayRef<Value> operandsRef,
                                            OperationFolder *folder) {
   SmallVector<Value, 4> operands(operandsRef.begin(), operandsRef.end());
   fullyComposeAffineMapAndOperands(&map, &operands);
   canonicalizeMapAndOperands(&map, &operands);
   return folder ? folder->create<AffineApplyOp>(b, loc, map, operands)

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

Value mlir::spirv::getBuiltinVariableValue(Operation *op,
 
   spirv::GlobalVariableOp varOp = getOrInsertBuiltinVariable(
       *parent->getRegion(0).begin(), op->getLoc(), builtin, builder);
   Value ptr = builder.create<spirv::AddressOfOp>(op->getLoc(), varOp);
   return builder.create<spirv::LoadOp>(op->getLoc(), ptr);
 }
 
 //===----------------------------------------------------------------------===//
+// Index calculation
+//===----------------------------------------------------------------------===//
+
+spirv::AccessChainOp mlir::spirv::getElementPtr(
+    SPIRVTypeConverter &typeConverter, MemRefType baseType, Value basePtr,
+    ArrayRef<Value> indices, Location loc, OpBuilder &builder) {
+  // Get base and offset of the MemRefType and verify they are static.
+  int64_t offset;
+  SmallVector<int64_t, 4> strides;
+  if (failed(getStridesAndOffset(baseType, strides, offset)) ||
+      llvm::is_contained(strides, MemRefType::getDynamicStrideOrOffset())) {
+    return nullptr;
+  }
+
+  auto indexType = typeConverter.getIndexType(builder.getContext());
+
+  Value ptrLoc = nullptr;
+  assert(indices.size() == strides.size());
+  for (auto index : enumerate(indices)) {
+    Value strideVal = builder.create<spirv::ConstantOp>(
+        loc, indexType, IntegerAttr::get(indexType, strides[index.index()]));
+    Value update = builder.create<spirv::IMulOp>(loc, strideVal, index.value());
+    ptrLoc =
+        (ptrLoc ? builder.create<spirv::IAddOp>(loc, ptrLoc, update).getResult()
+                : update);
+  }
+  SmallVector<Value, 2> linearizedIndices;
+  // Add a '0' at the start to index into the struct.
+  linearizedIndices.push_back(builder.create<spirv::ConstantOp>(
+      loc, indexType, IntegerAttr::get(indexType, 0)));
+  linearizedIndices.push_back(ptrLoc);
+  return builder.create<spirv::AccessChainOp>(loc, basePtr, linearizedIndices);
+}
+
+//===----------------------------------------------------------------------===//
 // Set ABI attributes for lowering entry functions.
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 mlir::spirv::setABIAttrs(FuncOp funcOp, spirv::EntryPointABIAttr entryPointInfo,
                          ArrayRef<spirv::InterfaceVarABIAttr> argABIInfo) {
   // Set the attributes for argument and the function.
   StringRef argABIAttrName = spirv::getInterfaceVarABIAttrName();

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

 static constexpr const char kEqualSemanticsAttrName[] = "equal_semantics";
 static constexpr const char kFnNameAttrName[] = "fn";
 static constexpr const char kGroupOperationAttrName[] = "group_operation";
 static constexpr const char kIndicesAttrName[] = "indices";
 static constexpr const char kInitializerAttrName[] = "initializer";
 static constexpr const char kInterfaceAttrName[] = "interface";
 static constexpr const char kMemoryScopeAttrName[] = "memory_scope";
 static constexpr const char kSemanticsAttrName[] = "semantics";
-static constexpr const char kSpecConstAttrName[] = "spec_const";
 static constexpr const char kSpecIdAttrName[] = "spec_id";
 static constexpr const char kTypeAttrName[] = "type";
 static constexpr const char kUnequalSemanticsAttrName[] = "unequal_semantics";
 static constexpr const char kValueAttrName[] = "value";
 static constexpr const char kValuesAttrName[] = "values";
-static constexpr const char kVariableAttrName[] = "variable";
 
 //===----------------------------------------------------------------------===//
 // Common utility functions
 //===----------------------------------------------------------------------===//
 
 /// Returns true if the given op is a function-like op or nested in a
 /// function-like op without a module-like op in the middle.
 static bool isNestedInFunctionLikeOp(Operation *op) {
void spirv::SelectionOp::addMergeBlock() {
   auto *mergeBlock = new Block();
   body().push_back(mergeBlock);
   OpBuilder builder(mergeBlock);
 
   // Add a spv._merge op into the merge block.
   builder.create<spirv::MergeOp>(getLoc());
 }
 
+spirv::SelectionOp spirv::SelectionOp::createIfThen(
+    Location loc, Value condition,
+    function_ref<void(OpBuilder *builder)> thenBody, OpBuilder *builder) {
+  auto selectionControl = builder->getI32IntegerAttr(
+      static_cast<uint32_t>(spirv::SelectionControl::None));
+  auto selectionOp = builder->create<spirv::SelectionOp>(loc, selectionControl);
+
+  selectionOp.addMergeBlock();
+  Block *mergeBlock = selectionOp.getMergeBlock();
+  Block *thenBlock = nullptr;
+
+  // Build the "then" block.
+  {
+    OpBuilder::InsertionGuard guard(*builder);
+    thenBlock = builder->createBlock(mergeBlock);
+    thenBody(builder);
+    builder->create<spirv::BranchOp>(loc, mergeBlock);
+  }
+
+  // Build the header block.
+  {
+    OpBuilder::InsertionGuard guard(*builder);
+    builder->createBlock(thenBlock);
+    builder->create<spirv::BranchConditionalOp>(
+        loc, condition, thenBlock,
+        /*trueArguments=*/ArrayRef<Value>(), mergeBlock,
+        /*falseArguments=*/ArrayRef<Value>());
+  }
+
+  return selectionOp;
+}
+
 namespace {
 // Blocks from the given `spv.selection` operation must satisfy the following
 // layout:
 //
 //       +-----------------------------------------------+
 //       | header block                                  |
 //       | spv.BranchConditionalOp %cond, ^case0, ^case1 |
 //       +-----------------------------------------------+

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

 //===- TargetAndABI.cpp - SPIR-V target and ABI utilities -----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/TargetAndABI.h"
 #include "mlir/Dialect/SPIRV/SPIRVTypes.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 
 namespace mlir {
 #include "mlir/Dialect/SPIRV/TargetAndABI.cpp.inc"
 }
 

 }
 
 spirv::TargetEnvAttr spirv::lookupTargetEnvOrDefault(Operation *op) {
   if (auto attr = op->getAttrOfType<spirv::TargetEnvAttr>(
           spirv::getTargetEnvAttrName()))
     return attr;
   return getDefaultTargetEnv(op->getContext());
 }
+
+DenseIntElementsAttr spirv::lookupLocalWorkGroupSize(Operation *op) {
+  while (op && !op->hasTrait<OpTrait::FunctionLike>())
+    op = op->getParentOp();
+  if (!op)
+    return {};
+
+  if (auto attr = op->getAttrOfType<spirv::EntryPointABIAttr>(
+          spirv::getEntryPointABIAttrName()))
+    return attr.local_size();
+
+  return {};
+}

mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir

+// RUN: mlir-opt -split-input-file -convert-linalg-to-spirv -canonicalize -verify-diagnostics %s -o - | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// Single workgroup reduction
+//===----------------------------------------------------------------------===//
+
+#single_workgroup_reduction_trait = {
+  args_in = 1,
+  args_out = 1,
+  iterator_types = ["reduction"],
+  indexing_maps = [
+    affine_map<(i) -> (i)>,
+    affine_map<(i) -> (0)>
+  ]
+}
+
+module attributes {
+  spv.target_env = {
+    version = 3 : i32,
+    extensions = [],
+    capabilities = [1: i32, 63: i32] // Shader, GroupNonUniformArithmetic
+  }
+} {
+
+// CHECK:      spv.globalVariable
+// CHECK-SAME: built_in("LocalInvocationId")
+
+// CHECK:      func @single_workgroup_reduction
+// CHECK-SAME: (%[[INPUT:.+]]: !spv.ptr{{.+}}, %[[OUTPUT:.+]]: !spv.ptr{{.+}})
+
+// CHECK:        %[[ZERO:.+]] = spv.constant 0 : i32
+// CHECK:        %[[ID:.+]] = spv.Load "Input" %{{.+}} : vector<3xi32>
+// CHECK:        %[[X:.+]] = spv.CompositeExtract %[[ID]][0 : i32]
+
+// CHECK:        %[[INPTR:.+]] = spv.AccessChain %[[INPUT]][%[[ZERO]], %[[X]]]
+// CHECK:        %[[VAL:.+]] = spv.Load "StorageBuffer" %[[INPTR]] : i32
+// CHECK:        %[[ADD:.+]] = spv.GroupNonUniformIAdd "Subgroup" "Reduce" %[[VAL]] : i32
+
+// CHECK:        %[[OUTPTR:.+]] = spv.AccessChain %[[OUTPUT]][%[[ZERO]], %[[ZERO]]]
+// CHECK:        %[[ELECT:.+]] = spv.GroupNonUniformElect "Subgroup" : i1
+
+// CHECK:        spv.selection {
+// CHECK:          spv.BranchConditional %[[ELECT]], ^bb1, ^bb2
+// CHECK:        ^bb1:
+// CHECK:          spv.AtomicIAdd "Device" "AcquireRelease" %[[OUTPTR]], %[[ADD]]
+// CHECK:          spv.Branch ^bb2
+// CHECK:        ^bb2:
+// CHECK:          spv._merge
+// CHECK:        }
+// CHECK:        spv.Return
+
+func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>) attributes {
+  spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}
+} {
+  linalg.generic #single_workgroup_reduction_trait %input, %output {
+    ^bb(%in: i32, %out: i32):
+      %sum = addi %in, %out : i32
+      linalg.yield %sum : i32
+  } : memref<16xi32>, memref<1xi32>
+  spv.Return
+}
+}
+
+// -----
+
+// Missing shader entry point ABI
+
+#single_workgroup_reduction_trait = {
+  args_in = 1,
+  args_out = 1,
+  iterator_types = ["reduction"],
+  indexing_maps = [
+    affine_map<(i) -> (i)>,
+    affine_map<(i) -> (0)>
+  ]
+}
+
+module attributes {
+  spv.target_env = {
+    version = 3 : i32,
+    extensions = [],
+    capabilities = [1: i32, 63: i32] // Shader, GroupNonUniformArithmetic
+  }
+} {
+func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>) {
+  // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
+  linalg.generic #single_workgroup_reduction_trait %input, %output {
+    ^bb(%in: i32, %out: i32):
+      %sum = addi %in, %out : i32
+      linalg.yield %sum : i32
+  } : memref<16xi32>, memref<1xi32>
+  return
+}
+}
+
+// -----
+
+// Mismatch between shader entry point ABI and input memref shape
+
+#single_workgroup_reduction_trait = {
+  args_in = 1,
+  args_out = 1,
+  iterator_types = ["reduction"],
+  indexing_maps = [
+    affine_map<(i) -> (i)>,
+    affine_map<(i) -> (0)>
+  ]
+}
+
+module attributes {
+  spv.target_env = {
+    version = 3 : i32,
+    extensions = [],
+    capabilities = [1: i32, 63: i32] // Shader, GroupNonUniformArithmetic
+  }
+} {
+func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>) attributes {
+  spv.entry_point_abi = {local_size = dense<[32, 1, 1]>: vector<3xi32>}
+} {
+  // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
+  linalg.generic #single_workgroup_reduction_trait %input, %output {
+    ^bb(%in: i32, %out: i32):
+      %sum = addi %in, %out : i32
+      linalg.yield %sum : i32
+  } : memref<16xi32>, memref<1xi32>
+  spv.Return
+}
+}
+
+// -----
+
+// Unsupported multi-dimension input memref
+
+#single_workgroup_reduction_trait = {
+  args_in = 1,
+  args_out = 1,
+  iterator_types = ["parallel", "reduction"],
+  indexing_maps = [
+    affine_map<(i, j) -> (i, j)>,
+    affine_map<(i, j) -> (i)>
+  ]
+}
+
+module attributes {
+  spv.target_env = {
+    version = 3 : i32,
+    extensions = [],
+    capabilities = [1: i32, 63: i32] // Shader, GroupNonUniformArithmetic
+  }
+} {
+func @single_workgroup_reduction(%input: memref<16x8xi32>, %output: memref<16xi32>) attributes {
+  spv.entry_point_abi = {local_size = dense<[16, 8, 1]>: vector<3xi32>}
+} {
+  // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
+  linalg.generic #single_workgroup_reduction_trait %input, %output {
+    ^bb(%in: i32, %out: i32):
+      %sum = addi %in, %out : i32
+      linalg.yield %sum : i32
+  } : memref<16x8xi32>, memref<16xi32>
+  spv.Return
+}
+}

mlir/tools/mlir-opt/CMakeLists.txt

set(LIBS
   MLIROpenMP
   MLIROptMain
   MLIRParser
   MLIRPass
   MLIRQuantizerTransforms
   MLIRQuantOps
   MLIRROCDLIR
   MLIRSPIRV
+  MLIRLinalgToSPIRVTransforms
   MLIRStandardToSPIRVTransforms
   MLIRSPIRVTestPasses
   MLIRSPIRVTransforms
   MLIRStandardOps
   MLIRStandardToLLVM
   MLIRTransforms
   MLIRTestDialect
   MLIRTestIR