Adding hal.interface.workgroup.* ops. (#4149)
diff --git a/iree/compiler/Conversion/LinalgToLLVM/ConvertToLLVM.cpp b/iree/compiler/Conversion/LinalgToLLVM/ConvertToLLVM.cpp index e5b7fe5..935300f 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/ConvertToLLVM.cpp +++ b/iree/compiler/Conversion/LinalgToLLVM/ConvertToLLVM.cpp
@@ -132,6 +132,34 @@ } }; +template <typename Op, int ArgIndex> +class ConvertWorkgroupInfoOpPattern : public ConvertToLLVMPattern { + public: + explicit ConvertWorkgroupInfoOpPattern(MLIRContext *context, + LLVMTypeConverter &typeConverter) + : ConvertToLLVMPattern(Op::getOperationName(), context, typeConverter) {} + + LogicalResult matchAndRewrite( + Operation *op, ArrayRef<Value> operands, + ConversionPatternRewriter &rewriter) const override { + auto newFuncOp = cast<LLVM::LLVMFuncOp>(rewriter.getBlock()->getParentOp()); + auto xyzTy = + LLVM::LLVMType::getInt32Ty(rewriter.getContext()).getPointerTo(); + auto xyzArgument = newFuncOp.getArgument(ArgIndex); + auto dimIndex = rewriter.createOrFold<LLVM::ConstantOp>( + op->getLoc(), LLVM::LLVMType::getInt64Ty(rewriter.getContext()), + op->getAttrOfType<IntegerAttr>("dimension")); + auto dimPtr = rewriter.createOrFold<LLVM::GEPOp>( + op->getLoc(), xyzTy, xyzArgument, ValueRange{dimIndex}); + auto dimValue = rewriter.createOrFold<LLVM::LoadOp>(op->getLoc(), dimPtr); + auto dimValueCasted = rewriter.createOrFold<LLVM::ZExtOp>( + op->getLoc(), typeConverter->convertType(op->getResult(0).getType()), + dimValue); + rewriter.replaceOp(op, dimValueCasted); + return success(); + } +}; + /// Returns true if `aOp` has a desciptor (set, binding) pair smaller than /// `bOp`. Note that this ignores the offset. bool operator<(IREE::HAL::InterfaceBindingOp aOp, @@ -168,17 +196,14 @@ // Get interface buffers from all the blocks. SmallVector<IREE::PlaceholderOp, 8> bufferOps; SmallVector<IREE::HAL::InterfaceLoadConstantOp, 8> loadOps; - SmallVector<IREE::WorkgroupIdOp, 3> workgroupIdOps; for (Block &block : funcOp.getBlocks()) { for (Operation &op : block) { - if (auto phOp = dyn_cast<IREE::PlaceholderOp>(op)) + if (auto phOp = dyn_cast<IREE::PlaceholderOp>(op)) { bufferOps.push_back(phOp); - if (auto phOp = dyn_cast<IREE::HAL::InterfaceLoadConstantOp>(op)) { + } else if (auto phOp = + dyn_cast<IREE::HAL::InterfaceLoadConstantOp>(op)) { loadOps.push_back(phOp); } - if (auto threadIdOp = dyn_cast<IREE::WorkgroupIdOp>(op)) { - workgroupIdOps.push_back(threadIdOp); - } } } @@ -233,14 +258,12 @@ auto packedBuffersArgsTy = LLVM::LLVMType::getInt8PtrTy(context).getPointerTo(); auto pushConstantArgTy = LLVM::LLVMType::getInt32Ty(context).getPointerTo(); - auto threadIdXTy = LLVM::LLVMType::getInt32Ty(context); - auto threadIdYTy = LLVM::LLVMType::getInt32Ty(context); - auto threadIdZTy = LLVM::LLVMType::getInt32Ty(context); + auto xyzTy = LLVM::LLVMType::getInt32Ty(context).getPointerTo(); signatureConverter.addInputs(packedBuffersArgsTy); signatureConverter.addInputs(pushConstantArgTy); - signatureConverter.addInputs(threadIdXTy); - signatureConverter.addInputs(threadIdYTy); - signatureConverter.addInputs(threadIdZTy); + signatureConverter.addInputs(xyzTy); // workgroup_id + signatureConverter.addInputs(xyzTy); // workgroup_count + signatureConverter.addInputs(xyzTy); // workgroup_size Location loc = funcOp.getLoc(); @@ -310,26 +333,6 @@ rewriter.replaceOp(loadOp, dimConstantCasted); } - // Lower iree.workgroup_idd ops to get indices from function arugments. - for (auto workgroupCoordOp : workgroupIdOps) { - auto attr = workgroupCoordOp.getAttrOfType<StringAttr>("dimension"); - int argIndex = -1; - if (attr.getValue().str() == "x") { - argIndex = 2; - } else if (attr.getValue().str() == "y") { - argIndex = 3; - } else if (attr.getValue().str() == "z") { - argIndex = 4; - } else { - return rewriter.notifyMatchFailure( - funcOp, - "Unable to map to workgroup coordinate : " + attr.getValue().str()); - } - Value threadXIndex = builder.create<LLVM::ZExtOp>( - loc, typeConverter->convertType(workgroupCoordOp.getType()), - newFuncOp.getArgument(argIndex)); - rewriter.replaceOp(workgroupCoordOp, threadXIndex); - } rewriter.eraseOp(funcOp); return success(); } @@ -389,6 +392,7 @@ populateVectorToLLVMMatrixConversionPatterns(converter, patterns); populateVectorToLLVMConversionPatterns(converter, patterns); populateLinalgToLLVMConversionPatterns(converter, patterns, &getContext()); + // The following patterns resolves dynamic shapes by substituting tie_shape // ops with an updated memref descriptors and replacing RankDimOp with // actual index loaded from memref<?xi32> that holds all dynamic shapes push @@ -396,6 +400,13 @@ patterns.insert<ConvertFuncWithHALInterface, ConvertRankedDimPattern, ConvertTieShapePattern, RemoveMakeRankedShape, RemoveInterfaceOpPattern>(&getContext(), converter); + + patterns.insert< + ConvertWorkgroupInfoOpPattern<IREE::HAL::InterfaceWorkgroupIDOp, 2>, + ConvertWorkgroupInfoOpPattern<IREE::HAL::InterfaceWorkgroupCountOp, 3>, + ConvertWorkgroupInfoOpPattern<IREE::HAL::InterfaceWorkgroupSizeOp, 4>>( + &getContext(), converter); + LLVMConversionTarget target(getContext()); target.addLegalOp<ModuleOp, ModuleTerminatorOp>();
diff --git a/iree/compiler/Conversion/LinalgToLLVM/LinalgTileAndDistributeOnTensorsPass.cpp b/iree/compiler/Conversion/LinalgToLLVM/LinalgTileAndDistributeOnTensorsPass.cpp index ff6a971..d3bdb3a 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/LinalgTileAndDistributeOnTensorsPass.cpp +++ b/iree/compiler/Conversion/LinalgToLLVM/LinalgTileAndDistributeOnTensorsPass.cpp
@@ -14,8 +14,8 @@ #include "iree/compiler/Conversion/CodegenUtils/MarkerUtils.h" #include "iree/compiler/Conversion/CodegenUtils/MatmulCodegenStrategy.h" -#include "iree/compiler/Dialect/IREE/IR/IREEDialect.h" -#include "iree/compiler/Dialect/IREE/IR/IREEOps.h" +#include "iree/compiler/Dialect/HAL/IR/HALDialect.h" +#include "iree/compiler/Dialect/HAL/IR/HALOps.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" #include "mlir/IR/Builders.h" #include "mlir/IR/MLIRContext.h" @@ -31,7 +31,7 @@ : public PassWrapper<LinalgTileAndDistributeOnTensorsPass, OperationPass<ModuleOp>> { void getDependentDialects(DialectRegistry ®istry) const override { - registry.insert<linalg::LinalgDialect, IREEDialect, AffineDialect, + registry.insert<linalg::LinalgDialect, IREE::HAL::HALDialect, AffineDialect, scf::SCFDialect>(); } LinalgTileAndDistributeOnTensorsPass() = default; @@ -45,16 +45,6 @@ llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated}; }; -static std::pair<Value, Value> buildWorkgroupOpPair(OpBuilder &b, - StringRef dim) { - Type indexType = b.getIndexType(); - StringAttr attr = b.getStringAttr(dim); - return {b.create<IREE::WorkgroupIdOp>(b.getInsertionPoint()->getLoc(), - indexType, attr), - b.create<IREE::WorkgroupSizeOp>(b.getInsertionPoint()->getLoc(), - indexType, attr)}; -} - // Rewrite pattern to ensure only ops with tensor semantics are tiled. struct TileAndDistributeOnTensorsPattern : public linalg::LinalgBaseTilingPattern { @@ -86,14 +76,16 @@ // range [0, WorkgroupSizeOp). static linalg::LinalgLoopDistributionOptions workgroupDistributionOptions = { [](OpBuilder &builder, Location loc, ArrayRef<Range> parallelLoopRanges) { - // TODO: drop magic names. - std::array<StringRef, 3> dimStrs{"x", "y", "z"}; - size_t numParallelDims = parallelLoopRanges.size(); - SmallVector<linalg::ProcInfo, 2> procInfo(numParallelDims); + auto numParallelDims = parallelLoopRanges.size(); + SmallVector<linalg::ProcInfo, 3> procInfo(numParallelDims); for (size_t dim = 0; dim < std::min(numParallelDims, static_cast<size_t>(3)); ++dim) { - auto p = buildWorkgroupOpPair(builder, dimStrs[dim]); - procInfo[dim] = {p.first, p.second}; + procInfo[numParallelDims - dim - 1] = { + builder.createOrFold<IREE::HAL::InterfaceWorkgroupIDOp>( + loc, builder.getIndexType(), APInt(64, dim)), + builder.createOrFold<IREE::HAL::InterfaceWorkgroupCountOp>( + loc, builder.getIndexType(), APInt(64, dim)), + }; } return procInfo; },
diff --git a/iree/compiler/Conversion/LinalgToLLVM/LinalgTileAndDistributePass.cpp b/iree/compiler/Conversion/LinalgToLLVM/LinalgTileAndDistributePass.cpp index cd0cdd2..9470e1c 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/LinalgTileAndDistributePass.cpp +++ b/iree/compiler/Conversion/LinalgToLLVM/LinalgTileAndDistributePass.cpp
@@ -19,8 +19,8 @@ #include "iree/compiler/Conversion/Common/Attributes.h" #include "iree/compiler/Conversion/Common/Transforms.h" #include "iree/compiler/Conversion/LinalgToLLVM/KernelDispatch.h" -#include "iree/compiler/Dialect/IREE/IR/IREEDialect.h" -#include "iree/compiler/Dialect/IREE/IR/IREEOps.h" +#include "iree/compiler/Dialect/HAL/IR/HALDialect.h" +#include "iree/compiler/Dialect/HAL/IR/HALOps.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" #include "mlir/IR/Builders.h" #include "mlir/IR/Matchers.h" @@ -37,7 +37,7 @@ struct LinalgTileAndDistributePass : public PassWrapper<LinalgTileAndDistributePass, OperationPass<ModuleOp>> { void getDependentDialects(DialectRegistry ®istry) const override { - registry.insert<linalg::LinalgDialect, IREEDialect, AffineDialect, + registry.insert<linalg::LinalgDialect, IREE::HAL::HALDialect, AffineDialect, scf::SCFDialect>(); } LinalgTileAndDistributePass() = default; @@ -155,16 +155,16 @@ static linalg::LinalgLoopDistributionOptions workgroupDistributionOptions = { [](OpBuilder &builder, Location loc, ArrayRef<Range> parallelLoopRanges) { - Type indexType = builder.getIndexType(); auto numParallelDims = parallelLoopRanges.size(); - SmallVector<linalg::ProcInfo, 2> procInfo(numParallelDims); - for (int dim = 0; dim < numParallelDims; ++dim) { - std::array<StringRef, 3> dimAttr{"x", "y", "z"}; - StringAttr attr = - builder.getStringAttr(dimAttr[std::min<unsigned>(dim, 3)]); + SmallVector<linalg::ProcInfo, 3> procInfo(numParallelDims); + for (size_t dim = 0; + dim < std::min(numParallelDims, static_cast<size_t>(3)); ++dim) { procInfo[numParallelDims - dim - 1] = { - builder.create<IREE::WorkgroupIdOp>(loc, indexType, attr), - builder.create<IREE::WorkgroupSizeOp>(loc, indexType, attr)}; + builder.createOrFold<IREE::HAL::InterfaceWorkgroupIDOp>( + loc, builder.getIndexType(), APInt(64, dim)), + builder.createOrFold<IREE::HAL::InterfaceWorkgroupCountOp>( + loc, builder.getIndexType(), APInt(64, dim)), + }; } return procInfo; },
diff --git a/iree/compiler/Conversion/LinalgToLLVM/Passes.h b/iree/compiler/Conversion/LinalgToLLVM/Passes.h index 42036b6..c8c3096 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/Passes.h +++ b/iree/compiler/Conversion/LinalgToLLVM/Passes.h
@@ -28,10 +28,10 @@ /// order. std::unique_ptr<FunctionPass> createPlanConvLoopOrderPass(); -/// Distributes linalg ops among iree.workgroup logical threads. +/// Distributes linalg ops among hal.interface.workgroup logical threads. std::unique_ptr<OperationPass<ModuleOp>> createLinalgTileAndDistributePass(); -/// Vectorizes linalg ops executed in the same iree.workgroup. +/// Vectorizes linalg ops executed in the same hal.interface.workgroup. std::unique_ptr<FunctionPass> createLinalgTileAndVectorizeWorkgroupsPass(); std::unique_ptr<OperationPass<ModuleOp>>
diff --git a/iree/compiler/Conversion/LinalgToLLVM/test/convert_to_llvm.mlir b/iree/compiler/Conversion/LinalgToLLVM/test/convert_to_llvm.mlir index fcae26a..5f244ea 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/test/convert_to_llvm.mlir +++ b/iree/compiler/Conversion/LinalgToLLVM/test/convert_to_llvm.mlir
@@ -14,7 +14,7 @@ hal.interface @legacy_io attributes {push_constants = 2 : i32, sym_visibility = "private"} { hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read" } -// CHECK: llvm.func @convert_dynamic_shape(%[[ARG0:.+]]: !llvm.ptr<ptr<i8>>, %[[ARG1:.+]]: !llvm.ptr<i32>, %[[THREAD_X_ID:.+]]: !llvm.i32, %[[THREAD_Y_ID:.+]]: !llvm.i32, %[[THREAD_Z_ID:.+]]: !llvm.i32) +// CHECK: llvm.func @convert_dynamic_shape(%[[ARG0:.+]]: !llvm.ptr<ptr<i8>>, %[[ARG1:.+]]: !llvm.ptr<i32>, %[[WORKGROUP_ID:.+]]: !llvm.ptr<i32>, %[[WORKGROUP_COUNT:.+]]: !llvm.ptr<i32>, %[[WORKGROUP_SIZE:.+]]: !llvm.ptr<i32>) // CHECK: %[[PACKED_ARGS_PTR:.+]] = llvm.bitcast %[[ARG0]] : !llvm.ptr<ptr<i8>> to !llvm.ptr<struct<(ptr<float>)>> // CHECK: %[[PACKED_ARGS:.+]] = llvm.load %[[PACKED_ARGS_PTR]] : !llvm.ptr<struct<(ptr<float>)>> // CHECK: %[[MEMREF0_DATA_PTR:.+]] = llvm.extractvalue %[[PACKED_ARGS]][0] : !llvm.struct<(ptr<float>)> @@ -53,7 +53,7 @@ hal.interface @legacy_io2 attributes {push_constants = 1 : i32, sym_visibility = "private"} { hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read" } -// CHECK: llvm.func @convert_dynamic_shape2(%[[ARG0:.+]]: !llvm.ptr<ptr<i8>>, %[[ARG1:.+]]: !llvm.ptr<i32>, %[[THREAD_X_ID:.+]]: !llvm.i32, %[[THREAD_Y_ID:.+]]: !llvm.i32, %[[THREAD_Z_ID:.+]]: !llvm.i32) +// CHECK: llvm.func @convert_dynamic_shape2(%[[ARG0:.+]]: !llvm.ptr<ptr<i8>>, %[[ARG1:.+]]: !llvm.ptr<i32>, %[[WORKGROUP_ID:.+]]: !llvm.ptr<i32>, %[[WORKGROUP_COUNT:.+]]: !llvm.ptr<i32>, %[[WORKGROUP_SIZE:.+]]: !llvm.ptr<i32>) // CHECK: %[[PACKED_ARGS_PTR:.+]] = llvm.bitcast %[[ARG0]] : !llvm.ptr<ptr<i8>> to !llvm.ptr<struct<(ptr<float>)>> // CHECK: %[[PACKED_ARGS:.+]] = llvm.load %[[PACKED_ARGS_PTR]] : !llvm.ptr<struct<(ptr<float>)>> // CHECK: %[[MEMREF0_DATA_PTR:.+]] = llvm.extractvalue %[[PACKED_ARGS]][0] : !llvm.struct<(ptr<float>)> @@ -86,9 +86,9 @@ // CHECK_LABEL: @distribute_lookup func @distribute_lookup() -> f32 { %0 = iree.placeholder for "interface buffer" {binding = @legacy_io3::@arg0} : memref<2x2x2xf32> - %1 = iree.workgroup_id {dimension = "x"} : index - %2 = iree.workgroup_id {dimension = "y"} : index - %3 = iree.workgroup_id {dimension = "z"} : index + %1 = hal.interface.workgroup.id[0] : index + %2 = hal.interface.workgroup.id[1] : index + %3 = hal.interface.workgroup.id[2] : index %4 = load %0[%1, %2, %3] : memref<2x2x2xf32> return %4 : f32 }
diff --git a/iree/compiler/Conversion/LinalgToLLVM/test/linalg_rewrite_destructive_updates.mlir b/iree/compiler/Conversion/LinalgToLLVM/test/linalg_rewrite_destructive_updates.mlir index 9360101..52b55fa 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/test/linalg_rewrite_destructive_updates.mlir +++ b/iree/compiler/Conversion/LinalgToLLVM/test/linalg_rewrite_destructive_updates.mlir
@@ -13,8 +13,8 @@ %2 = hal.interface.load.tensor @legacy_io::@TENSOR_INIT, offset = %c0 {operand_result_index = 2 : i32} : tensor<2x4xf32> - %3 = iree.workgroup_id {dimension = "x"} : index - %4 = iree.workgroup_id {dimension = "y"} : index + %3 = hal.interface.workgroup.id[0] : index + %4 = hal.interface.workgroup.id[1] : index // Test step %{{[0-9a-z]}} { to ensure yield has been folded away. // CHECK: scf.for %[[I:.*]] = {{.*}} step %{{[0-9a-z]+}} { // CHECK: scf.for %[[J:.*]] = {{.*}} step %{{[0-9a-z]+}} { @@ -83,8 +83,8 @@ %2 = hal.interface.load.tensor @legacy_io::@TENSOR_RHS, offset = %c0 {operand_result_index = 2 : i32} : tensor<3x4xf32> - %4 = iree.workgroup_id {dimension = "x"} : index - %5 = iree.workgroup_id {dimension = "y"} : index + %4 = hal.interface.workgroup.id[0] : index + %5 = hal.interface.workgroup.id[1] : index // Test step %{{[0-9a-z]}} { to ensure yield has been folded away. // CHECK: scf.for %[[I:.*]] = {{.*}} step %{{[0-9a-z]+}} { // CHECK: scf.for %[[J:.*]] = {{.*}} step %{{[0-9a-z]+}} { @@ -168,8 +168,8 @@ linalg.yield %arg0 : f32 } -> tensor<2x4xf32> - %4 = iree.workgroup_id {dimension = "x"} : index - %5 = iree.workgroup_id {dimension = "y"} : index + %4 = hal.interface.workgroup.id[0] : index + %5 = hal.interface.workgroup.id[1] : index // Test step %{{[0-9a-z]}} { to ensure yield has been folded away. // CHECK: scf.for %[[I:.*]] = {{.*}} step %{{[0-9a-z]+}} {
diff --git a/iree/compiler/Conversion/LinalgToLLVM/test/matmul_vectorization.mlir b/iree/compiler/Conversion/LinalgToLLVM/test/matmul_vectorization.mlir index d474f29..7ce18b9 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/test/matmul_vectorization.mlir +++ b/iree/compiler/Conversion/LinalgToLLVM/test/matmul_vectorization.mlir
@@ -7,8 +7,8 @@ // CHECK: #[[MAP0:map.*]] = affine_map<()[s0] -> (s0 * 64)> // CHECK-LABEL: func @matmul_128x128x128 // CHECK-SAME: (%[[ARG0:.+]]: memref<128x128xf32>, %[[ARG1:.+]]: memref<128x128xf32>, %[[ARG2:.+]]: memref<128x128xf32>) -// CHECK-DaG: %[[WORKGROUP_TILE_X:.+]] = iree.workgroup_id {dimension = "x"} : index -// CHECK-DAG: %[[WORKGROUP_TILE_Y:.+]] = iree.workgroup_id {dimension = "y"} : index +// CHECK-DaG: %[[WORKGROUP_TILE_X:.+]] = hal.interface.workgroup.id[0] : index +// CHECK-DAG: %[[WORKGROUP_TILE_Y:.+]] = hal.interface.workgroup.id[1] : index // CHECK-DAG: %[[START:.+]] = constant 0 // CHECK-DAG: %[[WORGKROUP_SIZE:.+]] = constant 64 // CHECK-DAG: %[[VECTOR_SIZE:.+]] = constant 4
diff --git a/iree/compiler/Conversion/LinalgToLLVM/test/tile_and_distribute.mlir b/iree/compiler/Conversion/LinalgToLLVM/test/tile_and_distribute.mlir index 35f816b..1eb024e 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/test/tile_and_distribute.mlir +++ b/iree/compiler/Conversion/LinalgToLLVM/test/tile_and_distribute.mlir
@@ -13,17 +13,17 @@ // CHECK-DAG: %[[CONST_0:.+]] = constant 0 : index // CHECK-DAG: %[[CONST_1:.+]] = constant 1 : index // CHECK-DAG: %[[DIM_K:.+]] = dim %[[LHS]], %[[CONST_1]] -// CHECK-DAG: %[[THREAD_X_ID:.+]] = iree.workgroup_id {dimension = "x"} : index -// CHECK-DAG: %[[THREAD_Y_ID:.+]] = iree.workgroup_id {dimension = "y"} : index +// CHECK-DAG: %[[THREAD_X_ID:.+]] = hal.interface.workgroup.id[0] : index +// CHECK-DAG: %[[THREAD_Y_ID:.+]] = hal.interface.workgroup.id[1] : index // CHECK: scf.for %[[K:.+]] = %[[CONST_0]] to %[[DIM_K]] // CHECK: %[[I:.+]] = affine.apply #[[MAP0]]()[%[[THREAD_Y_ID]]] // CHECK: %[[DIM_I:.+]] = dim %[[LHS]], %[[CONST_0]] // CHECK: %[[I_OFFSET:.+]] = affine.min #[[MAP1]]()[%[[THREAD_Y_ID]], %[[DIM_I]]] -// CHECK: %[[LHS_SUBVIEW:.+]] = subview %[[LHS]][%[[I]], %[[K]]] [%[[I_OFFSET]], 1] [1, 1] +// CHECK: %[[LHS_SUBVIEW:.+]] = subview %[[LHS]][%[[I]], %[[K]]] [%[[I_OFFSET]], 1] [1, 1] // CHECK: %[[J:.+]] = affine.apply #[[MAP3]]()[%[[THREAD_X_ID]]] -// CHECK: %[[DIM_J:.+]] = dim %[[RHS]], %[[CONST_1]] +// CHECK: %[[DIM_J:.+]] = dim %[[RHS]], %[[CONST_1]] // CHECK: %[[J_OFFSET:.+]] = affine.min #[[MAP4]]()[%[[THREAD_X_ID]], %[[DIM_J]]] -// CHECK: %[[RHS_SUBVIEW:.+]] = subview %[[RHS]][%[[K]], %[[J]]] [1, %[[J_OFFSET]]] [1, 1] +// CHECK: %[[RHS_SUBVIEW:.+]] = subview %[[RHS]][%[[K]], %[[J]]] [1, %[[J_OFFSET]]] [1, 1] // CHECK: %[[DIM_I:.+]] = dim %[[RESULT]], %[[CONST_0]] // CHECK: %[[DIM_I_OFFSET:.+]] = affine.min #[[MAP1]]()[%[[THREAD_Y_ID]], %[[DIM_I]]] // CHECK: %[[DIM_J:.+]] = dim %[[RESULT]], %[[CONST_1]] @@ -45,12 +45,12 @@ // CHECK-DAG: %[[CONST_0:.+]] = constant 0 : index // CHECK-DAG: %[[CONST_4:.+]] = constant 4 : index // CHECK-DAG: %[[CONST_1:.+]] = constant 1 : index -// CHECK-DAG: %[[THREAD_X_ID:.+]] = iree.workgroup_id {dimension = "x"} : index -// CHECK-DAG: %[[THREAD_Y_ID:.+]] = iree.workgroup_id {dimension = "y"} : index -// CHECK: scf.for %[[K:.+]] = %[[CONST_0]] to %[[CONST_4]] step %[[CONST_1]] +// CHECK-DAG: %[[THREAD_X_ID:.+]] = hal.interface.workgroup.id[0] : index +// CHECK-DAG: %[[THREAD_Y_ID:.+]] = hal.interface.workgroup.id[1] : index +// CHECK: scf.for %[[K:.+]] = %[[CONST_0]] to %[[CONST_4]] step %[[CONST_1]] // CHECK: %[[I:.+]] = affine.apply #[[MAP0]]()[%[[THREAD_Y_ID]]] // CHECK: %[[LHS_SUBVIEW:.+]] = subview %[[LHS]][%[[I]], %[[K]]] [2, 1] [1, 1] : memref<16x4xf32> to memref<2x1xf32, #[[MAP1]]> // CHECK: %[[J:.+]] = affine.apply #[[MAP2]]()[%[[THREAD_X_ID]]] // CHECK: %[[RHS_SUBVIEW:.+]] = subview %[[RHS]][%[[K]], %[[J]]] [1, 4] [1, 1] : memref<4x8xf32> to memref<1x4xf32, #[[MAP3]]> // CHECK: %[[RESULT_SUBVIEW:.+]] = subview %[[RESULT]][%[[I]], %[[J]]] [2, 4] [1, 1] : memref<16x8xf32> to memref<2x4xf32, #[[MAP3]]> -// CHECK: linalg.matmul {__internal_linalg_transform__ = "workgroup"} ins(%[[LHS_SUBVIEW]], %[[RHS_SUBVIEW]] : memref<2x1xf32, #[[MAP1]]>, memref<1x4xf32, #[[MAP3]]>) outs(%6 : memref<2x4xf32, #[[MAP3]]>) +// CHECK: linalg.matmul {__internal_linalg_transform__ = "workgroup"} ins(%[[LHS_SUBVIEW]], %[[RHS_SUBVIEW]] : memref<2x1xf32, #[[MAP1]]>, memref<1x4xf32, #[[MAP3]]>) outs(%4 : memref<2x4xf32, #[[MAP3]]>)
diff --git a/iree/compiler/Conversion/LinalgToLLVM/test/tile_and_distribute_on_tensors.mlir b/iree/compiler/Conversion/LinalgToLLVM/test/tile_and_distribute_on_tensors.mlir index b869782..6f35c1a 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/test/tile_and_distribute_on_tensors.mlir +++ b/iree/compiler/Conversion/LinalgToLLVM/test/tile_and_distribute_on_tensors.mlir
@@ -13,14 +13,14 @@ // CHECK-DAG: %[[C1:.*]] = constant 1 : index // CHECK-DAG: %[[C2:.*]] = constant 2 : index // CHECK-DAG: %[[C4:.*]] = constant 4 : index - // CHECK-DAG: %[[bix:.*]] = iree.workgroup_id {dimension = "x"} : index - // CHECK-DAG: %[[bdx:.*]] = iree.workgoup_size {dimension = "x"} : index - // CHECK-DAG: %[[biy:.*]] = iree.workgroup_id {dimension = "y"} : index - // CHECK-DAG: %[[bdy:.*]] = iree.workgoup_size {dimension = "y"} : index - // CHECK: %{{.*}} = scf.for %[[I:.*]] = %[[bix]] to %[[C2]] step %[[bdx]] iter_args(%arg1 = %2) -> (tensor<2x4xf32>) { - // CHECK-NEXT: %[[biy_scaled:.*]] = muli %[[biy]], %[[C2]] : index - // CHECK-NEXT: %[[bdy_scaled:.*]] = muli %[[bdy]], %[[C2]] : index - // CHECK-NEXT: %{{.*}} = scf.for %[[J:.*]] = %[[biy_scaled]] to %[[C4]] step %[[bdy_scaled]] iter_args(%arg3 = %arg1) -> (tensor<2x4xf32>) { + // CHECK-DAG: %[[bix:.*]] = hal.interface.workgroup.id[0] : index + // CHECK-DAG: %[[bdx:.*]] = hal.interface.workgroup.count[0] : index + // CHECK-DAG: %[[biy:.*]] = hal.interface.workgroup.id[1] : index + // CHECK-DAG: %[[bdy:.*]] = hal.interface.workgroup.count[1] : index + // CHECK: %{{.*}} = scf.for %[[I:.*]] = %[[biy]] to %[[C2]] step %[[bdy]] iter_args(%arg1 = %2) -> (tensor<2x4xf32>) { + // CHECK-NEXT: %[[bix_scaled:.*]] = muli %[[bix]], %[[C2]] : index + // CHECK-NEXT: %[[bdx_scaled:.*]] = muli %[[bdx]], %[[C2]] : index + // CHECK-NEXT: %{{.*}} = scf.for %[[J:.*]] = %[[bix_scaled]] to %[[C4]] step %[[bdx_scaled]] iter_args(%arg3 = %arg1) -> (tensor<2x4xf32>) { // CHECK-NEXT: subtensor %{{.*}}[%[[I]], 0] [1, 3] [1, 1] : tensor<2x3xf32> to tensor<1x3xf32> // // Canonicalizations not yet powerful enough here.
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/FoldGPUProcessorIDUses.cpp b/iree/compiler/Conversion/LinalgToSPIRV/FoldGPUProcessorIDUses.cpp index 1dda99d..5aee585 100644 --- a/iree/compiler/Conversion/LinalgToSPIRV/FoldGPUProcessorIDUses.cpp +++ b/iree/compiler/Conversion/LinalgToSPIRV/FoldGPUProcessorIDUses.cpp
@@ -73,6 +73,8 @@ case mlir::AffineExprKind::DimId: case mlir::AffineExprKind::SymbolId: return true; + default: + llvm_unreachable("unhandled affine expr kind"); } }
diff --git a/iree/compiler/Dialect/HAL/Conversion/IREEToHAL/ConvertIREEToHAL.cpp b/iree/compiler/Dialect/HAL/Conversion/IREEToHAL/ConvertIREEToHAL.cpp index dfa3657..f3493e1 100644 --- a/iree/compiler/Dialect/HAL/Conversion/IREEToHAL/ConvertIREEToHAL.cpp +++ b/iree/compiler/Dialect/HAL/Conversion/IREEToHAL/ConvertIREEToHAL.cpp
@@ -24,6 +24,7 @@ namespace iree_compiler { namespace { + class DynamicShapeConstantOpConversion : public OpConversionPattern<IREE::DynamicShapeConstantOp> { public:
diff --git a/iree/compiler/Dialect/HAL/IR/HALOps.cpp b/iree/compiler/Dialect/HAL/IR/HALOps.cpp index e6cfd96..c4a5b85 100644 --- a/iree/compiler/Dialect/HAL/IR/HALOps.cpp +++ b/iree/compiler/Dialect/HAL/IR/HALOps.cpp
@@ -18,6 +18,7 @@ #include "iree/compiler/Dialect/IREE/IR/IREETypes.h" #include "llvm/ADT/STLExtras.h" #include "llvm/Support/SMLoc.h" +#include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/Builders.h" #include "mlir/IR/OpImplementation.h" @@ -1580,6 +1581,44 @@ } //===----------------------------------------------------------------------===// +// hal.interface.workgroup.* +//===----------------------------------------------------------------------===// + +static void getAsmResultNamesForInterfaceWorkgroupOp( + StringRef prefix, const APInt &dimension, Value result, + function_ref<void(Value, StringRef)> setNameFn) { + switch (dimension.getZExtValue()) { + case 0: + setNameFn(result, (prefix + "x").str()); + return; + case 1: + setNameFn(result, (prefix + "y").str()); + return; + case 2: + setNameFn(result, (prefix + "z").str()); + return; + } +} + +void InterfaceWorkgroupIDOp::getAsmResultNames( + function_ref<void(Value, StringRef)> setNameFn) { + getAsmResultNamesForInterfaceWorkgroupOp("workgroup_id_", dimension(), + result(), setNameFn); +} + +void InterfaceWorkgroupCountOp::getAsmResultNames( + function_ref<void(Value, StringRef)> setNameFn) { + getAsmResultNamesForInterfaceWorkgroupOp("workgroup_count_", dimension(), + result(), setNameFn); +} + +void InterfaceWorkgroupSizeOp::getAsmResultNames( + function_ref<void(Value, StringRef)> setNameFn) { + getAsmResultNamesForInterfaceWorkgroupOp("workgroup_size_", dimension(), + result(), setNameFn); +} + +//===----------------------------------------------------------------------===// // hal.interface.load.tensor //===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Dialect/HAL/IR/HALOps.td b/iree/compiler/Dialect/HAL/IR/HALOps.td index 67d40de..c214e8d 100644 --- a/iree/compiler/Dialect/HAL/IR/HALOps.td +++ b/iree/compiler/Dialect/HAL/IR/HALOps.td
@@ -1856,7 +1856,8 @@ StrAttr:$sym_name, HAL_OrdinalAttr:$ordinal, FlatSymbolRefAttr:$interface, - TypeAttr:$signature + TypeAttr:$signature, + OptionalAttr<HAL_WorkgroupSizeAttr>:$workgroup_size ); } @@ -2083,6 +2084,80 @@ ); } +def HAL_InterfaceWorkgroupIDOp : HAL_PureOp<"interface.workgroup.id", [ + DeclareOpInterfaceMethods<OpAsmOpInterface>, + ]> { + let summary = [{returns the index of the current workgroup in the grid}]; + let description = [{ + The global workgroup ID of the current tile in the range of + `[0, hal.interface.workgroup.count)` along each XYZ dimension. + + Corresponds to the `WorkgroupId` SPIR-V built-in and the `blockIdx` CUDA + built-in variable. + + ```mlir + %x = hal.interface.workgroup.id[0] : index + %y = hal.interface.workgroup.id[1] : index + %z = hal.interface.workgroup.id[2] : index + ``` + }]; + + let arguments = (ins IndexAttr:$dimension); + let results = (outs HAL_Dim:$result); + + let assemblyFormat = "`[` $dimension `]` attr-dict `:` type($result)"; +} + +def HAL_InterfaceWorkgroupCountOp : HAL_PureOp<"interface.workgroup.count", [ + DeclareOpInterfaceMethods<OpAsmOpInterface>, + ]> { + let summary = [{returns the total workgroup count of the grid}]; + let description = [{ + The total number of workgroups along each dimension in the dispatch grid. + Matches what was passed to the `hal.command_buffer.dispatch` command (or + what was indirectly specified). + + Corresponds to the `NumWorkgroups` SPIR-V built-in and the `gridDim` CUDA + built-in variable. + + ```mlir + %x = hal.interface.workgroup.count[0] : index + %y = hal.interface.workgroup.count[1] : index + %z = hal.interface.workgroup.count[2] : index + ``` + }]; + + let arguments = (ins IndexAttr:$dimension); + let results = (outs HAL_Dim:$result); + + let assemblyFormat = "`[` $dimension `]` attr-dict `:` type($result)"; +} + +def HAL_InterfaceWorkgroupSizeOp : HAL_PureOp<"interface.workgroup.size", [ + DeclareOpInterfaceMethods<OpAsmOpInterface>, + ]> { + let summary = [{returns the size of each workgroup in invocations}]; + let description = [{ + The number of local invocations within the current workgroup along each + dimension. Depending on backend this may map to the SIMT thread count or + inner loop nest parameters. + + Corresponds to the `WorkgroupSize` SPIR-V built-in and the `blockDim` CUDA + built-in variable. + + ```mlir + %x = hal.interface.workgroup.size[0] : index + %y = hal.interface.workgroup.size[1] : index + %z = hal.interface.workgroup.size[2] : index + ``` + }]; + + let arguments = (ins IndexAttr:$dimension); + let results = (outs HAL_Dim:$result); + + let assemblyFormat = "`[` $dimension `]` attr-dict `:` type($result)"; +} + def HAL_InterfaceLoadConstantOp : HAL_PureOp<"interface.load.constant"> { let summary = [{loads a constant value from the interface constant block}]; let description = [{
diff --git a/iree/compiler/Dialect/HAL/IR/test/executable_ops.mlir b/iree/compiler/Dialect/HAL/IR/test/executable_ops.mlir index a94b474..8dbca61 100644 --- a/iree/compiler/Dialect/HAL/IR/test/executable_ops.mlir +++ b/iree/compiler/Dialect/HAL/IR/test/executable_ops.mlir
@@ -2,21 +2,6 @@ // RUN: iree-opt -allow-unregistered-dialect -split-input-file %s | iree-opt -allow-unregistered-dialect -split-input-file | IreeFileCheck %s -// CHECK-LABEL: @interface_io -func @interface_io() { - %c16 = constant 16 : index - // CHECK: %[[ARG0:.+]] = hal.interface.load.tensor @interface::@s0b0, offset = %c16 : tensor<4xf32> - %arg0 = hal.interface.load.tensor @interface::@s0b0, offset=%c16 : tensor<4xf32> - // CHECK-NEXT: %[[TEMP:.+]] = mhlo.add %[[ARG0]], %[[ARG0]] - %0 = mhlo.add %arg0, %arg0 : tensor<4xf32> - %c32 = constant 32 : index - // CHECK: hal.interface.store.tensor %[[TEMP]], @interface::@s0b1, offset = %c32 : tensor<4xf32> - hal.interface.store.tensor %0, @interface::@s0b1, offset=%c32 : tensor<4xf32> - return -} - -// ----- - // CHECK-LABEL: @ex hal.executable @ex { // CHECK: hal.executable.target @backend, filter="backend" @@ -94,22 +79,3 @@ %executable_layout = hal.executable_layout.create %arg0, set_layouts = [%arg1], push_constants = 1 : !hal.executable_layout return } - -// ----- - -// CHECK-LABEL: @interface_io -func @interface_io() { - %c16 = constant 16 : index - // CHECK: %[[ARG0:.+]] = hal.interface.load.tensor.tile @interface::@s0b0, base_offset = %c16 - // CHECK-SAME: offsets = [0], sizes = [4], strides = [1] : tensor<4xf32> - %arg0 = hal.interface.load.tensor.tile @interface::@s0b0, base_offset = %c16, - offsets = [0], sizes = [4], strides = [1]: tensor<4xf32> - // CHECK-NEXT: %[[TEMP:.+]] = mhlo.add %[[ARG0]], %[[ARG0]] - %0 = mhlo.add %arg0, %arg0 : tensor<4xf32> - %c32 = constant 32 : index - // CHECK: hal.interface.store.tensor.tile %[[TEMP]], @interface::@s0b1, base_offset = %c32 - // CHECK-SAME: offsets = [4], sizes = [7], strides = [1] : tensor<4xf32> - hal.interface.store.tensor.tile %0, @interface::@s0b1, base_offset = %c32, - offsets = [4], sizes = [7], strides = [1]: tensor<4xf32> - return -}
diff --git a/iree/compiler/Dialect/HAL/IR/test/interface_ops.mlir b/iree/compiler/Dialect/HAL/IR/test/interface_ops.mlir new file mode 100644 index 0000000..8a1ce16 --- /dev/null +++ b/iree/compiler/Dialect/HAL/IR/test/interface_ops.mlir
@@ -0,0 +1,46 @@ +// RUN: iree-opt -allow-unregistered-dialect -split-input-file %s | iree-opt -allow-unregistered-dialect -split-input-file | IreeFileCheck %s + +// CHECK-LABEL: @interface_workgroup_info +func @interface_workgroup_info() { + // CHECK: %workgroup_id_x = hal.interface.workgroup.id[0] : index + %0 = hal.interface.workgroup.id[0] : index + // CHECK: %workgroup_count_y = hal.interface.workgroup.count[1] : index + %1 = hal.interface.workgroup.count[1] : index + // CHECK: %workgroup_size_z = hal.interface.workgroup.size[2] : index + %2 = hal.interface.workgroup.size[2] : index + return +} + +// ----- + +// CHECK-LABEL: @interface_io_tensors +func @interface_io_tensors() { + %c16 = constant 16 : index + // CHECK: %[[ARG0:.+]] = hal.interface.load.tensor @interface::@s0b0, offset = %c16 : tensor<4xf32> + %arg0 = hal.interface.load.tensor @interface::@s0b0, offset=%c16 : tensor<4xf32> + // CHECK-NEXT: %[[TEMP:.+]] = mhlo.add %[[ARG0]], %[[ARG0]] + %0 = mhlo.add %arg0, %arg0 : tensor<4xf32> + %c32 = constant 32 : index + // CHECK: hal.interface.store.tensor %[[TEMP]], @interface::@s0b1, offset = %c32 : tensor<4xf32> + hal.interface.store.tensor %0, @interface::@s0b1, offset=%c32 : tensor<4xf32> + return +} + +// ----- + +// CHECK-LABEL: @interface_io_tiles +func @interface_io_tiles() { + %c16 = constant 16 : index + // CHECK: %[[ARG0:.+]] = hal.interface.load.tensor.tile @interface::@s0b0, base_offset = %c16 + // CHECK-SAME: offsets = [0], sizes = [4], strides = [1] : tensor<4xf32> + %arg0 = hal.interface.load.tensor.tile @interface::@s0b0, base_offset = %c16, + offsets = [0], sizes = [4], strides = [1]: tensor<4xf32> + // CHECK-NEXT: %[[TEMP:.+]] = mhlo.add %[[ARG0]], %[[ARG0]] + %0 = mhlo.add %arg0, %arg0 : tensor<4xf32> + %c32 = constant 32 : index + // CHECK: hal.interface.store.tensor.tile %[[TEMP]], @interface::@s0b1, base_offset = %c32 + // CHECK-SAME: offsets = [4], sizes = [7], strides = [1] : tensor<4xf32> + hal.interface.store.tensor.tile %0, @interface::@s0b1, base_offset = %c32, + offsets = [4], sizes = [7], strides = [1]: tensor<4xf32> + return +}
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMBaseTarget.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMBaseTarget.cpp index 93484e8..d640698 100644 --- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMBaseTarget.cpp +++ b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMBaseTarget.cpp
@@ -165,7 +165,7 @@ entryPointOp.getLoc(), entryPointOp.sym_nameAttr(), builder.getI32IntegerAttr(nextEntryPointOrdinal++), builder.getSymbolRefAttr(interfaceOpForExecutable.getName()), - entryPointOp.signatureAttr()); + entryPointOp.signatureAttr(), ArrayAttr{}); // Add to replacement table for fixing up dispatch calls referencing // this entry point.
diff --git a/iree/compiler/Dialect/HAL/Target/VMLA/VMLATarget.cpp b/iree/compiler/Dialect/HAL/Target/VMLA/VMLATarget.cpp index 2ce8686..12c5568 100644 --- a/iree/compiler/Dialect/HAL/Target/VMLA/VMLATarget.cpp +++ b/iree/compiler/Dialect/HAL/Target/VMLA/VMLATarget.cpp
@@ -169,7 +169,7 @@ entryPointOp.getLoc(), entryPointOp.sym_nameAttr(), builder.getI32IntegerAttr(nextEntryPointOrdinal++), builder.getSymbolRefAttr(interfaceOpForExecutable.getName()), - entryPointOp.signatureAttr()); + entryPointOp.signatureAttr(), ArrayAttr{}); // Add to replacement table for fixing up dispatch calls referencing // this entry point.
diff --git a/iree/compiler/Dialect/HAL/Transforms/BUILD b/iree/compiler/Dialect/HAL/Transforms/BUILD index d240022..59c702d 100644 --- a/iree/compiler/Dialect/HAL/Transforms/BUILD +++ b/iree/compiler/Dialect/HAL/Transforms/BUILD
@@ -31,6 +31,7 @@ "MemoizeDeviceQueries.cpp", "PackConstantPoolStorage.cpp", "Passes.cpp", + "PropagateConstantWorkgroupInfo.cpp", "PublicAbiGeneration.cpp", "ResolveEntryPointOrdinals.cpp", "SerializeExecutables.cpp",
diff --git a/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt b/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt index ed48d8a..5a0ec5e 100644 --- a/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt +++ b/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
@@ -30,6 +30,7 @@ "MemoizeDeviceQueries.cpp" "PackConstantPoolStorage.cpp" "Passes.cpp" + "PropagateConstantWorkgroupInfo.cpp" "PublicAbiGeneration.cpp" "ResolveEntryPointOrdinals.cpp" "SerializeExecutables.cpp"
diff --git a/iree/compiler/Dialect/HAL/Transforms/MaterializeInterfaces.cpp b/iree/compiler/Dialect/HAL/Transforms/MaterializeInterfaces.cpp index 01b7f87..f1fe326 100644 --- a/iree/compiler/Dialect/HAL/Transforms/MaterializeInterfaces.cpp +++ b/iree/compiler/Dialect/HAL/Transforms/MaterializeInterfaces.cpp
@@ -277,7 +277,7 @@ builder.getStringAttr(thunkFuncOp->getName()), builder.getI32IntegerAttr(nextOrdinal++), builder.getSymbolRefAttr(interfaceOp), - TypeAttr::get(sourceFuncOp.getType())); + TypeAttr::get(sourceFuncOp.getType()), ArrayAttr{}); } }
diff --git a/iree/compiler/Dialect/HAL/Transforms/Passes.cpp b/iree/compiler/Dialect/HAL/Transforms/Passes.cpp index 4c078d9..27a8762 100644 --- a/iree/compiler/Dialect/HAL/Transforms/Passes.cpp +++ b/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
@@ -62,6 +62,8 @@ passManager.addPass(createMaterializeInterfacesPass(targetOptions)); passManager.nest<ExecutableOp>().addNestedPass<ExecutableTargetOp>( + createPropagateConstantWorkgroupInfoPass()); + passManager.nest<ExecutableOp>().addNestedPass<ExecutableTargetOp>( createTranslateExecutablesPass(targetOptions)); // Convert supported input dialects (std, flow, etc) into the HAL dialect.
diff --git a/iree/compiler/Dialect/HAL/Transforms/Passes.h b/iree/compiler/Dialect/HAL/Transforms/Passes.h index a0ab331..d148aea 100644 --- a/iree/compiler/Dialect/HAL/Transforms/Passes.h +++ b/iree/compiler/Dialect/HAL/Transforms/Passes.h
@@ -75,6 +75,10 @@ std::unique_ptr<OperationPass<ModuleOp>> createMaterializeInterfacesPass( TargetOptions executableOptions); +// Propagates hal.interface.workload.* information when constant. +std::unique_ptr<OperationPass<IREE::HAL::ExecutableTargetOp>> +createPropagateConstantWorkgroupInfoPass(); + // Translates hal.executable.target ops via a nested translation pipeline. std::unique_ptr<OperationPass<IREE::HAL::ExecutableTargetOp>> createTranslateExecutablesPass(TargetOptions executableOptions);
diff --git a/iree/compiler/Dialect/HAL/Transforms/PropagateConstantWorkgroupInfo.cpp b/iree/compiler/Dialect/HAL/Transforms/PropagateConstantWorkgroupInfo.cpp new file mode 100644 index 0000000..9994970 --- /dev/null +++ b/iree/compiler/Dialect/HAL/Transforms/PropagateConstantWorkgroupInfo.cpp
@@ -0,0 +1,68 @@ +// Copyright 2020 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "iree/compiler/Dialect/HAL/IR/HALOps.h" +#include "iree/compiler/Dialect/HAL/Transforms/Passes.h" +#include "mlir/IR/SymbolTable.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace iree_compiler { +namespace IREE { +namespace HAL { + +class PropagateConstantWorkgroupInfoPass + : public PassWrapper<PropagateConstantWorkgroupInfoPass, + OperationPass<IREE::HAL::ExecutableTargetOp>> { + public: + void runOnOperation() override { + auto targetOp = getOperation(); + + SymbolTable targetSymbolTable(targetOp); + for (auto funcOp : targetOp.getInnerModule().getOps<FuncOp>()) { + auto entryPointOp = + targetSymbolTable.lookup<IREE::HAL::ExecutableEntryPointOp>( + funcOp.getName()); + if (!entryPointOp) continue; + if (!entryPointOp.workgroup_size().hasValue()) continue; + auto workgroupSizeAttr = entryPointOp.workgroup_sizeAttr(); + auto workgroupSizeOps = llvm::to_vector<4>( + funcOp.getOps<IREE::HAL::InterfaceWorkgroupSizeOp>()); + for (auto workgroupSizeOp : workgroupSizeOps) { + OpBuilder builder(workgroupSizeOp); + auto dimValue = builder.createOrFold<ConstantIndexOp>( + workgroupSizeOp.getLoc(), + workgroupSizeAttr[workgroupSizeOp.dimension().getZExtValue()] + .cast<IntegerAttr>() + .getInt()); + workgroupSizeOp.replaceAllUsesWith(dimValue); + workgroupSizeOp.erase(); + } + } + } +}; + +std::unique_ptr<OperationPass<IREE::HAL::ExecutableTargetOp>> +createPropagateConstantWorkgroupInfoPass() { + return std::make_unique<PropagateConstantWorkgroupInfoPass>(); +} + +static PassRegistration<PropagateConstantWorkgroupInfoPass> pass( + "iree-hal-propagate-constant-workgroup-info", + "Propagates constant hal.interface.workgroup.* queries when known"); + +} // namespace HAL +} // namespace IREE +} // namespace iree_compiler +} // namespace mlir
diff --git a/iree/compiler/Dialect/HAL/Transforms/test/propagate_constant_workgroup_info.mlir b/iree/compiler/Dialect/HAL/Transforms/test/propagate_constant_workgroup_info.mlir new file mode 100644 index 0000000..6efe132 --- /dev/null +++ b/iree/compiler/Dialect/HAL/Transforms/test/propagate_constant_workgroup_info.mlir
@@ -0,0 +1,26 @@ +// RUN: iree-opt -allow-unregistered-dialect -split-input-file -pass-pipeline='hal.executable(hal.executable.target(iree-hal-propagate-constant-workgroup-info))' %s | IreeFileCheck %s + +hal.executable @exe { + hal.interface @interface { + hal.interface.binding @s0b0, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @s0b1, set=0, binding=1, type="StorageBuffer", access="Read|Write" + } + hal.executable.target @target, filter="target" { + hal.executable.entry_point @entry attributes { + interface = @interface, + ordinal = 0 : i32, + signature = (tensor<4xf32>) -> tensor<4xf32>, + workgroup_size = [32 : index, 4 : index, 8 : index] + } + module { + // CHECK: func @entry() + func @entry() { + // CHECK-DAG: constant 32 : index + %workgroup_size_x = hal.interface.workgroup.size[0] : index + // CHECK-DAG: constant 4 : index + %workgroup_size_y = hal.interface.workgroup.size[1] : index + return + } + } + } +}
diff --git a/iree/compiler/Dialect/IREE/IR/IREEOps.td b/iree/compiler/Dialect/IREE/IR/IREEOps.td index 8fcb71d..184134d 100644 --- a/iree/compiler/Dialect/IREE/IR/IREEOps.td +++ b/iree/compiler/Dialect/IREE/IR/IREEOps.td
@@ -56,9 +56,12 @@ } //===----------------------------------------------------------------------===// -// Executable ABI +// Executable ABI DO NOT ADD THINGS HERE //===----------------------------------------------------------------------===// +// ***WARNING***: this is going away soon - please don't extend or add more +// things like this to this file. Ops should go to the dialect they relate to +// (HAL, in this case). def IREE_PlaceholderOp : IREE_Op<"placeholder", [MemoryEffects<[MemAlloc]>]> { let summary = "A placeholder op to feed a value/buffer into computation"; let description = [{ @@ -75,38 +78,6 @@ let assemblyFormat = [{ `for` $purpose attr-dict `:` type($output) }]; } -def IREE_WorkgroupIdOp : IREE_PureOp<"workgroup_id"> { - let summary = "Get grid index of an iree workgoup among a specific dimension."; - let description = [{ - IREE workgroups are logically distributed among a hypergrid, each point sampled - from the grid corresponds to a logical thread. For example in a 3d grid case - the op quries (x, y, z) dimensions: - - ```mlir - %0 = iree.workgroup_id {dimension = "x"} : index - %1 = iree.workgroup_id {dimension = "y"} : index - %2 = iree.workgroup_id {dimension = "z"} : index - ``` - }]; - - let arguments = (ins StrAttr:$dimension); - let results = (outs Index:$result); - - let assemblyFormat = "attr-dict `:` type($result)"; -} - -def IREE_WorkgroupSizeOp: IREE_PureOp<"workgoup_size"> { - let summary = "Get grid size of an iree thread among a specific dimension"; - let description = [{ - Get grid size of an iree workgroups among a specific dimension - }]; - let arguments = (ins StrAttr:$dimension); - let results = (outs Index:$result); - - let assemblyFormat = "attr-dict `:` type($result)"; -} - - //===----------------------------------------------------------------------===// // Compiler hints //===----------------------------------------------------------------------===//
diff --git a/iree/hal/dylib/dylib_executable.cc b/iree/hal/dylib/dylib_executable.cc index 9c93e6e..2aa7292 100644 --- a/iree/hal/dylib/dylib_executable.cc +++ b/iree/hal/dylib/dylib_executable.cc
@@ -188,6 +188,8 @@ IREE_TRACE(const char* entry_name = nullptr); + std::array<uint32_t, 3> workgroup_count; + std::array<uint32_t, 3> workgroup_size; void* entry_function = nullptr; std::array<void*, 32> args; std::array<uint32_t, 32> push_constants; @@ -203,6 +205,8 @@ } auto dispatch_state = make_ref<DyLibDispatchState>(); + dispatch_state->workgroup_count = params.workgroup_count; + dispatch_state->workgroup_size = params.workgroup_size; IREE_TRACE(dispatch_state->entry_name = entry_names_[params.entry_point]); dispatch_state->entry_function = entry_functions_[params.entry_point]; @@ -229,11 +233,12 @@ auto* dispatch_state = static_cast<DyLibDispatchState*>(state); IREE_TRACE_SCOPE_DYNAMIC(dispatch_state->entry_name); - auto entry_function = (void (*)(void**, uint32_t*, int32_t, int32_t, - int32_t))dispatch_state->entry_function; + auto entry_function = (void (*)(void**, uint32_t*, uint32_t*, uint32_t*, + uint32_t*))dispatch_state->entry_function; entry_function(dispatch_state->args.data(), - dispatch_state->push_constants.data(), workgroup_xyz[0], - workgroup_xyz[1], workgroup_xyz[2]); + dispatch_state->push_constants.data(), workgroup_xyz.data(), + dispatch_state->workgroup_count.data(), + dispatch_state->workgroup_size.data()); return OkStatus(); }
diff --git a/iree/hal/host/host_executable.h b/iree/hal/host/host_executable.h index cb2e6c9..9b7aa78 100644 --- a/iree/hal/host/host_executable.h +++ b/iree/hal/host/host_executable.h
@@ -46,6 +46,9 @@ // Total workgroup XYZ count for the grid. std::array<uint32_t, 3> workgroup_count; + // Size of each tile in the grid in local space. + std::array<uint32_t, 3> workgroup_size; + // Push constants populated by the command buffer. const PushConstantBlock* push_constants = nullptr;
diff --git a/iree/hal/llvmjit/llvmjit_executable.cc b/iree/hal/llvmjit/llvmjit_executable.cc index 5dd1a82..2db4bd8 100644 --- a/iree/hal/llvmjit/llvmjit_executable.cc +++ b/iree/hal/llvmjit/llvmjit_executable.cc
@@ -176,6 +176,8 @@ struct LLVMJITDispatchState : public HostExecutable::DispatchState { LLVMJITDispatchState() = default; + std::array<uint32_t, 3> workgroup_count; + std::array<uint32_t, 3> workgroup_size; llvm::JITEvaluatedSymbol symbol; llvm::SmallVector<void*, 4> args; llvm::SmallVector<int32_t, 4> push_constant; @@ -191,6 +193,8 @@ } auto dispatch_state = make_ref<LLVMJITDispatchState>(); + dispatch_state->workgroup_count = params.workgroup_count; + dispatch_state->workgroup_size = params.workgroup_size; dispatch_state->symbol = symbols_[params.entry_point]; for (size_t set = 0; set < params.set_bindings.size(); ++set) { @@ -218,10 +222,11 @@ IREE_TRACE_SCOPE0("LLVMJITExecutable::DispatchTile"); auto* dispatch_state = static_cast<LLVMJITDispatchState*>(state); - auto func_ptr = (void (*)(void**, int32_t*, int32_t, int32_t, - int32_t))dispatch_state->symbol.getAddress(); + auto func_ptr = (void (*)(void**, int32_t*, uint32_t*, uint32_t*, + uint32_t*))dispatch_state->symbol.getAddress(); func_ptr(dispatch_state->args.data(), dispatch_state->push_constant.data(), - workgroup_xyz[0], workgroup_xyz[1], workgroup_xyz[2]); + workgroup_xyz.data(), dispatch_state->workgroup_count.data(), + dispatch_state->workgroup_size.data()); return OkStatus(); }